diff --git a/api/services/tools_manage_service.py b/api/services/tools_manage_service.py index fa033b39f3..0144d7f381 100644 --- a/api/services/tools_manage_service.py +++ b/api/services/tools_manage_service.py @@ -1,6 +1,5 @@ import json -from flask import current_app from httpx import get from core.tools.entities.common_entities import I18nObject @@ -33,40 +32,18 @@ class ToolManageService: :return: the list of tool providers """ - result = [provider.to_dict() for provider in ToolManager.user_list_providers( + providers = ToolManager.user_list_providers( user_id, tenant_id - )] + ) - # add icon url prefix - for provider in result: - ToolManageService.repack_provider(provider) + # add icon + for provider in providers: + ToolTransformService.repack_provider(provider) + + result = [provider.to_dict() for provider in providers] return result - @staticmethod - def repack_provider(provider: dict): - """ - repack provider - - :param provider: the provider dict - """ - url_prefix = (current_app.config.get("CONSOLE_API_URL") - + "/console/api/workspaces/current/tool-provider/") - - if 'icon' in provider: - if provider['type'] == UserToolProvider.ProviderType.BUILTIN.value: - provider['icon'] = url_prefix + 'builtin/' + provider['name'] + '/icon' - elif provider['type'] == UserToolProvider.ProviderType.MODEL.value: - provider['icon'] = url_prefix + 'model/' + provider['name'] + '/icon' - elif provider['type'] == UserToolProvider.ProviderType.API.value: - try: - provider['icon'] = json.loads(provider['icon']) - except: - provider['icon'] = { - "background": "#252525", - "content": "\ud83d\ude01" - } - @staticmethod def list_builtin_tool_provider_tools( user_id: str, tenant_id: str, provider: str @@ -651,7 +628,10 @@ class ToolManageService: provider_controller=provider_controller, db_provider=find_provider(provider_controller.identity.name) ) - + + # add icon + ToolTransformService.repack_provider(user_builtin_provider) + tools = provider_controller.get_tools() for tool in tools: user_builtin_provider.tools.append(ToolTransformService.tool_to_user_tool( @@ -686,9 +666,13 @@ class ToolManageService: db_provider=provider ) + # add icon + ToolTransformService.repack_provider(user_provider) + tools = provider_controller.get_tools( user_id=user_id, tenant_id=tenant_id ) + for tool in tools: user_provider.tools.append(ToolTransformService.tool_to_user_tool( tenant_id=tenant_id, diff --git a/api/services/tools_transform_service.py b/api/services/tools_transform_service.py index 8db7db62e4..861eab73c5 100644 --- a/api/services/tools_transform_service.py +++ b/api/services/tools_transform_service.py @@ -1,5 +1,8 @@ +import json import logging -from typing import Optional +from typing import Optional, Union + +from flask import current_app from core.model_runtime.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ApiProviderAuthType, ToolParameter, ToolProviderCredentials @@ -14,6 +17,49 @@ from models.tools import ApiToolProvider, BuiltinToolProvider logger = logging.getLogger(__name__) class ToolTransformService: + @staticmethod + def get_tool_provider_icon_url(provider_type: str, provider_name: str, icon: str) -> Union[str, dict]: + """ + get tool provider icon url + """ + url_prefix = (current_app.config.get("CONSOLE_API_URL") + + "/console/api/workspaces/current/tool-provider/") + + if provider_type == UserToolProvider.ProviderType.BUILTIN.value: + return url_prefix + 'builtin/' + provider_name + '/icon' + elif provider_type == UserToolProvider.ProviderType.MODEL.value: + return url_prefix + 'model/' + provider_name + '/icon' + elif provider_type == UserToolProvider.ProviderType.API.value: + try: + return json.loads(icon) + except: + return { + "background": "#252525", + "content": "\ud83d\ude01" + } + + return '' + + @staticmethod + def repack_provider(provider: Union[dict, UserToolProvider]): + """ + repack provider + + :param provider: the provider dict + """ + if isinstance(provider, dict) and 'icon' in provider: + provider['icon'] = ToolTransformService.get_tool_provider_icon_url( + provider_type=provider['type'], + provider_name=provider['name'], + icon=provider['icon'] + ) + elif isinstance(provider, UserToolProvider): + provider.icon = ToolTransformService.get_tool_provider_icon_url( + provider_type=provider.type.value, + provider_name=provider.name, + icon=provider.icon + ) + @staticmethod def builtin_provider_to_user_provider( provider_controller: BuiltinToolProviderController,