tjdify/api/core/tools/tool_manager.py

import json
import logging
import mimetypes
from os import listdir, path
from typing import Any, Union

from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler
from core.entities.application_entities import AgentToolEntity
from core.model_runtime.entities.message_entities import PromptMessage
from core.provider_manager import ProviderManager
from core.tools.entities.common_entities import I18nObject
from core.tools.entities.constant import DEFAULT_PROVIDERS
from core.tools.entities.tool_entities import (
    ApiProviderAuthType,
    ToolInvokeMessage,
    ToolParameter,
    ToolProviderCredentials,
)
from core.tools.entities.user_entities import UserToolProvider
from core.tools.errors import ToolProviderNotFoundError
from core.tools.provider.api_tool_provider import ApiBasedToolProviderController
from core.tools.provider.app_tool_provider import AppBasedToolProviderEntity
from core.tools.provider.builtin._positions import BuiltinToolProviderSort
from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
from core.tools.provider.model_tool_provider import ModelToolProviderController
from core.tools.provider.tool_provider import ToolProviderController
from core.tools.tool.api_tool import ApiTool
from core.tools.tool.builtin_tool import BuiltinTool
from core.tools.tool.tool import Tool
from core.tools.utils.configuration import (
    ModelToolConfigurationManager,
    ToolConfigurationManager,
    ToolParameterConfigurationManager,
)
from core.tools.utils.encoder import serialize_base_model_dict
from core.utils.module_import_helper import load_single_subclass_from_source
from extensions.ext_database import db
from models.tools import ApiToolProvider, BuiltinToolProvider

logger = logging.getLogger(__name__)

_builtin_providers = {}
_builtin_tools_labels = {}

class ToolManager:
    @staticmethod
    def invoke(
        provider: str,
        tool_id: str,
        tool_name: str,
        tool_parameters: dict[str, Any],
        credentials: dict[str, Any],
        prompt_messages: list[PromptMessage],
    ) -> list[ToolInvokeMessage]:
        """
            invoke the assistant

            :param provider: the name of the provider
            :param tool_id: the id of the tool
            :param tool_name: the name of the tool, defined in `get_tools`
            :param tool_parameters: the parameters of the tool
            :param credentials: the credentials of the tool
            :param prompt_messages: the prompt messages that the tool can use

            :return: the messages that the tool wants to send to the user
        """
        provider_entity: ToolProviderController = None
        if provider == DEFAULT_PROVIDERS.API_BASED:
            provider_entity = ApiBasedToolProviderController()
        elif provider == DEFAULT_PROVIDERS.APP_BASED:
            provider_entity = AppBasedToolProviderEntity()

        if provider_entity is None:
            # fetch the provider from .provider.builtin
            provider_class = load_single_subclass_from_source(
                module_name=f'core.tools.provider.builtin.{provider}.{provider}',
                script_path=path.join(path.dirname(path.realpath(__file__)), 'builtin', provider, f'{provider}.py'),
                parent_type=ToolProviderController)
            provider_entity = provider_class()

        return provider_entity.invoke(tool_id, tool_name, tool_parameters, credentials, prompt_messages)
    
    @staticmethod
    def get_builtin_provider(provider: str) -> BuiltinToolProviderController:
        global _builtin_providers
        """
            get the builtin provider

            :param provider: the name of the provider
            :return: the provider
        """
        if len(_builtin_providers) == 0:
            # init the builtin providers
            ToolManager.list_builtin_providers()

        if provider not in _builtin_providers:
            raise ToolProviderNotFoundError(f'builtin provider {provider} not found')
        
        return _builtin_providers[provider]
    
    @staticmethod
    def get_builtin_tool(provider: str, tool_name: str) -> BuiltinTool:
        """
            get the builtin tool

            :param provider: the name of the provider
            :param tool_name: the name of the tool

            :return: the provider, the tool
        """
        provider_controller = ToolManager.get_builtin_provider(provider)
        tool = provider_controller.get_tool(tool_name)

        return tool
    
    @staticmethod
    def get_tool(provider_type: str, provider_id: str, tool_name: str, tenant_id: str = None) \
        -> Union[BuiltinTool, ApiTool]:
        """
            get the tool

            :param provider_type: the type of the provider
            :param provider_name: the name of the provider
            :param tool_name: the name of the tool

            :return: the tool
        """
        if provider_type == 'builtin':
            return ToolManager.get_builtin_tool(provider_id, tool_name)
        elif provider_type == 'api':
            if tenant_id is None:
                raise ValueError('tenant id is required for api provider')
            api_provider, _ = ToolManager.get_api_provider_controller(tenant_id, provider_id)
            return api_provider.get_tool(tool_name)
        elif provider_type == 'app':
            raise NotImplementedError('app provider not implemented')
        else:
            raise ToolProviderNotFoundError(f'provider type {provider_type} not found')
        
    @staticmethod
    def get_tool_runtime(provider_type: str, provider_name: str, tool_name: str, tenant_id: str, 
                         agent_callback: DifyAgentCallbackHandler = None) \
        -> Union[BuiltinTool, ApiTool]:
        """
            get the tool runtime

            :param provider_type: the type of the provider
            :param provider_name: the name of the provider
            :param tool_name: the name of the tool

            :return: the tool
        """
        if provider_type == 'builtin':
            builtin_tool = ToolManager.get_builtin_tool(provider_name, tool_name)

            # check if the builtin tool need credentials
            provider_controller = ToolManager.get_builtin_provider(provider_name)
            if not provider_controller.need_credentials:
                return builtin_tool.fork_tool_runtime(meta={
                    'tenant_id': tenant_id,
                    'credentials': {},
                }, agent_callback=agent_callback)

            # get credentials
            builtin_provider: BuiltinToolProvider = db.session.query(BuiltinToolProvider).filter(
                BuiltinToolProvider.tenant_id == tenant_id,
                BuiltinToolProvider.provider == provider_name,
            ).first()

            if builtin_provider is None:
                raise ToolProviderNotFoundError(f'builtin provider {provider_name} not found')
            
            # decrypt the credentials
            credentials = builtin_provider.credentials
            controller = ToolManager.get_builtin_provider(provider_name)
            tool_configuration = ToolConfigurationManager(tenant_id=tenant_id, provider_controller=controller)

            decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials)

            return builtin_tool.fork_tool_runtime(meta={
                'tenant_id': tenant_id,
                'credentials': decrypted_credentials,
                'runtime_parameters': {}
            }, agent_callback=agent_callback)
        
        elif provider_type == 'api':
            if tenant_id is None:
                raise ValueError('tenant id is required for api provider')
            
            api_provider, credentials = ToolManager.get_api_provider_controller(tenant_id, provider_name)

            # decrypt the credentials
            tool_configuration = ToolConfigurationManager(tenant_id=tenant_id, provider_controller=api_provider)
            decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials)

            return api_provider.get_tool(tool_name).fork_tool_runtime(meta={
                'tenant_id': tenant_id,
                'credentials': decrypted_credentials,
            })
        elif provider_type == 'model':
            if tenant_id is None:
                raise ValueError('tenant id is required for model provider')
            # get model provider
            model_provider = ToolManager.get_model_provider(tenant_id, provider_name)

            # get tool
            model_tool = model_provider.get_tool(tool_name)

            return model_tool.fork_tool_runtime(meta={
                'tenant_id': tenant_id,
                'credentials': model_tool.model_configuration['model_instance'].credentials
            })
        elif provider_type == 'app':
            raise NotImplementedError('app provider not implemented')
        else:
            raise ToolProviderNotFoundError(f'provider type {provider_type} not found')

    @staticmethod
    def get_agent_tool_runtime(tenant_id: str, agent_tool: AgentToolEntity, agent_callback: DifyAgentCallbackHandler) -> Tool:
        """
            get the agent tool runtime
        """
        tool_entity = ToolManager.get_tool_runtime(
            provider_type=agent_tool.provider_type, provider_name=agent_tool.provider_id, tool_name=agent_tool.tool_name, 
            tenant_id=tenant_id,
            agent_callback=agent_callback
        )
        runtime_parameters = {}
        parameters = tool_entity.get_all_runtime_parameters()
        for parameter in parameters:
            if parameter.form == ToolParameter.ToolParameterForm.FORM:
                # get tool parameter from form
                tool_parameter_config = agent_tool.tool_parameters.get(parameter.name)
                if not tool_parameter_config:
                    # get default value
                    tool_parameter_config = parameter.default
                    if not tool_parameter_config and parameter.required:
                        raise ValueError(f"tool parameter {parameter.name} not found in tool config")
                    
                if parameter.type == ToolParameter.ToolParameterType.SELECT:
                    # check if tool_parameter_config in options
                    options = list(map(lambda x: x.value, parameter.options))
                    if tool_parameter_config not in options:
                        raise ValueError(f"tool parameter {parameter.name} value {tool_parameter_config} not in options {options}")
                    
                # convert tool parameter config to correct type
                try:
                    if parameter.type == ToolParameter.ToolParameterType.NUMBER:
                        # check if tool parameter is integer
                        if isinstance(tool_parameter_config, int):
                            tool_parameter_config = tool_parameter_config
                        elif isinstance(tool_parameter_config, float):
                            tool_parameter_config = tool_parameter_config
                        elif isinstance(tool_parameter_config, str):
                            if '.' in tool_parameter_config:
                                tool_parameter_config = float(tool_parameter_config)
                            else:
                                tool_parameter_config = int(tool_parameter_config)
                    elif parameter.type == ToolParameter.ToolParameterType.BOOLEAN:
                        tool_parameter_config = bool(tool_parameter_config)
                    elif parameter.type not in [ToolParameter.ToolParameterType.SELECT, ToolParameter.ToolParameterType.STRING]:
                        tool_parameter_config = str(tool_parameter_config)
                    elif parameter.type == ToolParameter.ToolParameterType:
                        tool_parameter_config = str(tool_parameter_config)
                except Exception as e:
                    raise ValueError(f"tool parameter {parameter.name} value {tool_parameter_config} is not correct type")
                
                # save tool parameter to tool entity memory
                runtime_parameters[parameter.name] = tool_parameter_config
        
        # decrypt runtime parameters
        encryption_manager = ToolParameterConfigurationManager(
            tenant_id=tenant_id,
            tool_runtime=tool_entity,
            provider_name=agent_tool.provider_id,
            provider_type=agent_tool.provider_type,
        )
        runtime_parameters = encryption_manager.decrypt_tool_parameters(runtime_parameters)

        tool_entity.runtime.runtime_parameters.update(runtime_parameters)
        return tool_entity

    @staticmethod
    def get_builtin_provider_icon(provider: str) -> tuple[str, str]:
        """
            get the absolute path of the icon of the builtin provider

            :param provider: the name of the provider

            :return: the absolute path of the icon, the mime type of the icon
        """
        # get provider
        provider_controller = ToolManager.get_builtin_provider(provider)

        absolute_path = path.join(path.dirname(path.realpath(__file__)), 'provider', 'builtin', provider, '_assets', provider_controller.identity.icon)
        # check if the icon exists
        if not path.exists(absolute_path):
            raise ToolProviderNotFoundError(f'builtin provider {provider} icon not found')
        
        # get the mime type
        mime_type, _ = mimetypes.guess_type(absolute_path)
        mime_type = mime_type or 'application/octet-stream'

        return absolute_path, mime_type

    @staticmethod
    def list_builtin_providers() -> list[BuiltinToolProviderController]:
        global _builtin_providers

        # use cache first
        if len(_builtin_providers) > 0:
            return list(_builtin_providers.values())
        
        builtin_providers: list[BuiltinToolProviderController] = []
        for provider in listdir(path.join(path.dirname(path.realpath(__file__)), 'provider', 'builtin')):
            if provider.startswith('__'):
                continue

            if path.isdir(path.join(path.dirname(path.realpath(__file__)), 'provider', 'builtin', provider)):
                if provider.startswith('__'):
                    continue

                # init provider
                provider_class = load_single_subclass_from_source(
                    module_name=f'core.tools.provider.builtin.{provider}.{provider}',
                    script_path=path.join(path.dirname(path.realpath(__file__)),
                                           'provider', 'builtin', provider, f'{provider}.py'),
                    parent_type=BuiltinToolProviderController)
                builtin_providers.append(provider_class())

        # cache the builtin providers
        for provider in builtin_providers:
            _builtin_providers[provider.identity.name] = provider
            for tool in provider.get_tools():
                _builtin_tools_labels[tool.identity.name] = tool.identity.label

        return builtin_providers
    
    @staticmethod
    def list_model_providers(tenant_id: str = None) -> list[ModelToolProviderController]:
        """
            list all the model providers

            :return: the list of the model providers
        """
        tenant_id = tenant_id or 'ffffffff-ffff-ffff-ffff-ffffffffffff'
        # get configurations
        model_configurations = ModelToolConfigurationManager.get_all_configuration()
        # get all providers
        provider_manager = ProviderManager()
        configurations = provider_manager.get_configurations(tenant_id).values()
        # get model providers
        model_providers: list[ModelToolProviderController] = []
        for configuration in configurations:
            # all the model tool should be configurated
            if configuration.provider.provider not in model_configurations:
                continue
            if not ModelToolProviderController.is_configuration_valid(configuration):
                continue
            model_providers.append(ModelToolProviderController.from_db(configuration))

        return model_providers
    
    @staticmethod
    def get_model_provider(tenant_id: str, provider_name: str) -> ModelToolProviderController:
        """
            get the model provider

            :param provider_name: the name of the provider

            :return: the provider
        """
        # get configurations
        provider_manager = ProviderManager()
        configurations = provider_manager.get_configurations(tenant_id)
        configuration = configurations.get(provider_name)
        if configuration is None:
            raise ToolProviderNotFoundError(f'model provider {provider_name} not found')
        
        return ModelToolProviderController.from_db(configuration)

    @staticmethod
    def get_tool_label(tool_name: str) -> Union[I18nObject, None]:
        """
            get the tool label

            :param tool_name: the name of the tool

            :return: the label of the tool
        """
        global _builtin_tools_labels
        if len(_builtin_tools_labels) == 0:
            # init the builtin providers
            ToolManager.list_builtin_providers()

        if tool_name not in _builtin_tools_labels:
            return None
        
        return _builtin_tools_labels[tool_name]
    
    @staticmethod
    def user_list_providers(
        user_id: str,
        tenant_id: str,
    ) -> list[UserToolProvider]:
        result_providers: dict[str, UserToolProvider] = {}
        # get builtin providers
        builtin_providers = ToolManager.list_builtin_providers()
        # append builtin providers
        for provider in builtin_providers:
            result_providers[provider.identity.name] = UserToolProvider(
                id=provider.identity.name,
                author=provider.identity.author,
                name=provider.identity.name,
                description=I18nObject(
                    en_US=provider.identity.description.en_US,
                    zh_Hans=provider.identity.description.zh_Hans,
                ),
                icon=provider.identity.icon,
                label=I18nObject(
                    en_US=provider.identity.label.en_US,
                    zh_Hans=provider.identity.label.zh_Hans,
                ),
                type=UserToolProvider.ProviderType.BUILTIN,
                team_credentials={},
                is_team_authorization=False,
            )

            # get credentials schema
            schema = provider.get_credentials_schema()
            for name, value in schema.items():
                result_providers[provider.identity.name].team_credentials[name] = \
                    ToolProviderCredentials.CredentialsType.default(value.type)

            # check if the provider need credentials
            if not provider.need_credentials:
                result_providers[provider.identity.name].is_team_authorization = True
                result_providers[provider.identity.name].allow_delete = False

        # get db builtin providers
        db_builtin_providers: list[BuiltinToolProvider] = db.session.query(BuiltinToolProvider). \
            filter(BuiltinToolProvider.tenant_id == tenant_id).all()
        
        for db_builtin_provider in db_builtin_providers:
            # add provider into providers
            credentials = db_builtin_provider.credentials
            provider_name = db_builtin_provider.provider
            result_providers[provider_name].is_team_authorization = True

            # package builtin tool provider controller
            controller = ToolManager.get_builtin_provider(provider_name)

            # init tool configuration
            tool_configuration = ToolConfigurationManager(tenant_id=tenant_id, provider_controller=controller)
            # decrypt the credentials and mask the credentials
            decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials=credentials)
            masked_credentials = tool_configuration.mask_tool_credentials(credentials=decrypted_credentials)

            result_providers[provider_name].team_credentials = masked_credentials

        # get model tool providers
        model_providers = ToolManager.list_model_providers(tenant_id=tenant_id)
        # append model providers
        for provider in model_providers:
            result_providers[f'model_provider.{provider.identity.name}'] = UserToolProvider(
                id=provider.identity.name,
                author=provider.identity.author,
                name=provider.identity.name,
                description=I18nObject(
                    en_US=provider.identity.description.en_US,
                    zh_Hans=provider.identity.description.zh_Hans,
                ),
                icon=provider.identity.icon,
                label=I18nObject(
                    en_US=provider.identity.label.en_US,
                    zh_Hans=provider.identity.label.zh_Hans,
                ),
                type=UserToolProvider.ProviderType.MODEL,
                team_credentials={},
                is_team_authorization=provider.is_active,
            )

        # get db api providers
        db_api_providers: list[ApiToolProvider] = db.session.query(ApiToolProvider). \
            filter(ApiToolProvider.tenant_id == tenant_id).all()
        
        for db_api_provider in db_api_providers:
            username = 'Anonymous'
            try:
                username = db_api_provider.user.name
            except Exception as e:
                logger.error(f'failed to get user name for api provider {db_api_provider.id}: {str(e)}')
            # add provider into providers
            credentials = db_api_provider.credentials
            provider_name = db_api_provider.name
            result_providers[provider_name] = UserToolProvider(
                id=db_api_provider.id,
                author=username,
                name=db_api_provider.name,
                description=I18nObject(
                    en_US=db_api_provider.description,
                    zh_Hans=db_api_provider.description,
                ),
                icon=db_api_provider.icon,
                label=I18nObject(
                    en_US=db_api_provider.name,
                    zh_Hans=db_api_provider.name,
                ),
                type=UserToolProvider.ProviderType.API,
                team_credentials={},
                is_team_authorization=True,
            )

            # package tool provider controller
            controller = ApiBasedToolProviderController.from_db(
                db_provider=db_api_provider,
                auth_type=ApiProviderAuthType.API_KEY if db_api_provider.credentials['auth_type'] == 'api_key' else ApiProviderAuthType.NONE
            )

            # init tool configuration
            tool_configuration = ToolConfigurationManager(tenant_id=tenant_id, provider_controller=controller)

            # decrypt the credentials and mask the credentials
            decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials=credentials)
            masked_credentials = tool_configuration.mask_tool_credentials(credentials=decrypted_credentials)

            result_providers[provider_name].team_credentials = masked_credentials

        return BuiltinToolProviderSort.sort(list(result_providers.values()))
    
    @staticmethod
    def get_api_provider_controller(tenant_id: str, provider_id: str) -> tuple[ApiBasedToolProviderController, dict[str, Any]]:
        """
            get the api provider

            :param provider_name: the name of the provider

            :return: the provider controller, the credentials
        """
        provider: ApiToolProvider = db.session.query(ApiToolProvider).filter(
            ApiToolProvider.id == provider_id,
            ApiToolProvider.tenant_id == tenant_id,
        ).first()

        if provider is None:
            raise ToolProviderNotFoundError(f'api provider {provider_id} not found')
        
        controller = ApiBasedToolProviderController.from_db(
            provider, ApiProviderAuthType.API_KEY if provider.credentials['auth_type'] == 'api_key' else ApiProviderAuthType.NONE
        )
        controller.load_bundled_tools(provider.tools)

        return controller, provider.credentials
    
    @staticmethod
    def user_get_api_provider(provider: str, tenant_id: str) -> dict:
        """
            get api provider
        """
        """
            get tool provider
        """
        provider: ApiToolProvider = db.session.query(ApiToolProvider).filter(
            ApiToolProvider.tenant_id == tenant_id,
            ApiToolProvider.name == provider,
        ).first()

        if provider is None:
            raise ValueError(f'you have not added provider {provider}')
        
        try:
            credentials = json.loads(provider.credentials_str) or {}
        except:
            credentials = {}

        # package tool provider controller
        controller = ApiBasedToolProviderController.from_db(
            provider, ApiProviderAuthType.API_KEY if credentials['auth_type'] == 'api_key' else ApiProviderAuthType.NONE
        )
        # init tool configuration
        tool_configuration = ToolConfigurationManager(tenant_id=tenant_id, provider_controller=controller)

        decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials)
        masked_credentials = tool_configuration.mask_tool_credentials(decrypted_credentials)

        try:
            icon = json.loads(provider.icon)
        except:
            icon = {
                "background": "#252525",
                "content": "\ud83d\ude01"
            }

        return json.loads(serialize_base_model_dict({
            'schema_type': provider.schema_type,
            'schema': provider.schema,
            'tools': provider.tools,
            'icon': icon,
            'description': provider.description,
            'credentials': masked_credentials,
            'privacy_policy': provider.privacy_policy
        }))