builtin_tools_manage_service.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267
  1. import json
  2. import logging
  3. from configs import dify_config
  4. from core.helper.position_helper import is_filtered
  5. from core.model_runtime.utils.encoders import jsonable_encoder
  6. from core.tools.entities.api_entities import UserTool, UserToolProvider
  7. from core.tools.errors import ToolNotFoundError, ToolProviderCredentialValidationError, ToolProviderNotFoundError
  8. from core.tools.provider.builtin._positions import BuiltinToolProviderSort
  9. from core.tools.provider.tool_provider import ToolProviderController
  10. from core.tools.tool_label_manager import ToolLabelManager
  11. from core.tools.tool_manager import ToolManager
  12. from core.tools.utils.configuration import ToolConfigurationManager
  13. from extensions.ext_database import db
  14. from models.tools import BuiltinToolProvider
  15. from services.tools.tools_transform_service import ToolTransformService
  16. logger = logging.getLogger(__name__)
  17. class BuiltinToolManageService:
  18. @staticmethod
  19. def list_builtin_tool_provider_tools(user_id: str, tenant_id: str, provider: str) -> list[UserTool]:
  20. """
  21. list builtin tool provider tools
  22. """
  23. provider_controller: ToolProviderController = ToolManager.get_builtin_provider(provider)
  24. tools = provider_controller.get_tools()
  25. tool_provider_configurations = ToolConfigurationManager(
  26. tenant_id=tenant_id,
  27. config=provider_controller.get_credentials_schema(),
  28. provider_type=provider_controller.provider_type.value,
  29. provider_identity=provider_controller.identity.name,
  30. )
  31. # check if user has added the provider
  32. builtin_provider: BuiltinToolProvider | None = (
  33. db.session.query(BuiltinToolProvider)
  34. .filter(
  35. BuiltinToolProvider.tenant_id == tenant_id,
  36. BuiltinToolProvider.provider == provider,
  37. )
  38. .first()
  39. )
  40. credentials = {}
  41. if builtin_provider is not None:
  42. # get credentials
  43. credentials = builtin_provider.credentials
  44. credentials = tool_provider_configurations.decrypt_tool_credentials(credentials)
  45. result = []
  46. for tool in tools:
  47. result.append(
  48. ToolTransformService.tool_to_user_tool(
  49. tool=tool,
  50. credentials=credentials,
  51. tenant_id=tenant_id,
  52. labels=ToolLabelManager.get_tool_labels(provider_controller),
  53. )
  54. )
  55. return result
  56. @staticmethod
  57. def list_builtin_provider_credentials_schema(provider_name):
  58. """
  59. list builtin provider credentials schema
  60. :return: the list of tool providers
  61. """
  62. provider = ToolManager.get_builtin_provider(provider_name)
  63. return jsonable_encoder([v for _, v in (provider.credentials_schema or {}).items()])
  64. @staticmethod
  65. def update_builtin_tool_provider(user_id: str, tenant_id: str, provider_name: str, credentials: dict):
  66. """
  67. update builtin tool provider
  68. """
  69. # get if the provider exists
  70. provider: BuiltinToolProvider | None = (
  71. db.session.query(BuiltinToolProvider)
  72. .filter(
  73. BuiltinToolProvider.tenant_id == tenant_id,
  74. BuiltinToolProvider.provider == provider_name,
  75. )
  76. .first()
  77. )
  78. try:
  79. # get provider
  80. provider_controller = ToolManager.get_builtin_provider(provider_name)
  81. if not provider_controller.need_credentials:
  82. raise ValueError(f"provider {provider_name} does not need credentials")
  83. tool_configuration = ToolConfigurationManager(
  84. tenant_id=tenant_id,
  85. config=provider_controller.get_credentials_schema(),
  86. provider_type=provider_controller.provider_type.value,
  87. provider_identity=provider_controller.identity.name,
  88. )
  89. # get original credentials if exists
  90. if provider is not None:
  91. original_credentials = tool_configuration.decrypt_tool_credentials(provider.credentials)
  92. masked_credentials = tool_configuration.mask_tool_credentials(original_credentials)
  93. # check if the credential has changed, save the original credential
  94. for name, value in credentials.items():
  95. if name in masked_credentials and value == masked_credentials[name]:
  96. credentials[name] = original_credentials[name]
  97. # validate credentials
  98. provider_controller.validate_credentials(credentials)
  99. # encrypt credentials
  100. credentials = tool_configuration.encrypt_tool_credentials(credentials)
  101. except (ToolProviderNotFoundError, ToolNotFoundError, ToolProviderCredentialValidationError) as e:
  102. raise ValueError(str(e))
  103. if provider is None:
  104. # create provider
  105. provider = BuiltinToolProvider(
  106. tenant_id=tenant_id,
  107. user_id=user_id,
  108. provider=provider_name,
  109. encrypted_credentials=json.dumps(credentials),
  110. )
  111. db.session.add(provider)
  112. db.session.commit()
  113. else:
  114. provider.encrypted_credentials = json.dumps(credentials)
  115. db.session.add(provider)
  116. db.session.commit()
  117. # delete cache
  118. tool_configuration.delete_tool_credentials_cache()
  119. return {"result": "success"}
  120. @staticmethod
  121. def get_builtin_tool_provider_credentials(user_id: str, tenant_id: str, provider: str):
  122. """
  123. get builtin tool provider credentials
  124. """
  125. provider_obj: BuiltinToolProvider | None = (
  126. db.session.query(BuiltinToolProvider)
  127. .filter(
  128. BuiltinToolProvider.tenant_id == tenant_id,
  129. BuiltinToolProvider.provider == provider,
  130. )
  131. .first()
  132. )
  133. if provider_obj is None:
  134. return {}
  135. provider_controller = ToolManager.get_builtin_provider(provider_obj.provider)
  136. tool_configuration = ToolConfigurationManager(
  137. tenant_id=tenant_id,
  138. config=provider_controller.get_credentials_schema(),
  139. provider_type=provider_controller.provider_type.value,
  140. provider_identity=provider_controller.identity.name,
  141. )
  142. credentials = tool_configuration.decrypt_tool_credentials(provider_obj.credentials)
  143. credentials = tool_configuration.mask_tool_credentials(credentials)
  144. return credentials
  145. @staticmethod
  146. def delete_builtin_tool_provider(user_id: str, tenant_id: str, provider_name: str):
  147. """
  148. delete tool provider
  149. """
  150. provider_obj: BuiltinToolProvider | None = (
  151. db.session.query(BuiltinToolProvider)
  152. .filter(
  153. BuiltinToolProvider.tenant_id == tenant_id,
  154. BuiltinToolProvider.provider == provider_name,
  155. )
  156. .first()
  157. )
  158. if provider_obj is None:
  159. raise ValueError(f"you have not added provider {provider_name}")
  160. db.session.delete(provider_obj)
  161. db.session.commit()
  162. # delete cache
  163. provider_controller = ToolManager.get_builtin_provider(provider_name)
  164. tool_configuration = ToolConfigurationManager(
  165. tenant_id=tenant_id,
  166. config=provider_controller.get_credentials_schema(),
  167. provider_type=provider_controller.provider_type.value,
  168. provider_identity=provider_controller.identity.name,
  169. )
  170. tool_configuration.delete_tool_credentials_cache()
  171. return {"result": "success"}
  172. @staticmethod
  173. def get_builtin_tool_provider_icon(provider: str):
  174. """
  175. get tool provider icon and it's mimetype
  176. """
  177. icon_path, mime_type = ToolManager.get_builtin_provider_icon(provider)
  178. with open(icon_path, "rb") as f:
  179. icon_bytes = f.read()
  180. return icon_bytes, mime_type
  181. @staticmethod
  182. def list_builtin_tools(user_id: str, tenant_id: str) -> list[UserToolProvider]:
  183. """
  184. list builtin tools
  185. """
  186. # get all builtin providers
  187. provider_controllers = ToolManager.list_builtin_providers()
  188. # get all user added providers
  189. db_providers: list[BuiltinToolProvider] = (
  190. db.session.query(BuiltinToolProvider).filter(BuiltinToolProvider.tenant_id == tenant_id).all() or []
  191. )
  192. # find provider
  193. find_provider = lambda provider: next(
  194. filter(lambda db_provider: db_provider.provider == provider, db_providers), None
  195. )
  196. result: list[UserToolProvider] = []
  197. for provider_controller in provider_controllers:
  198. try:
  199. # handle include, exclude
  200. if is_filtered(
  201. include_set=dify_config.POSITION_TOOL_INCLUDES_SET, # type: ignore
  202. exclude_set=dify_config.POSITION_TOOL_EXCLUDES_SET, # type: ignore
  203. data=provider_controller,
  204. name_func=lambda x: x.identity.name,
  205. ):
  206. continue
  207. # convert provider controller to user provider
  208. user_builtin_provider = ToolTransformService.builtin_provider_to_user_provider(
  209. provider_controller=provider_controller,
  210. db_provider=find_provider(provider_controller.identity.name),
  211. decrypt_credentials=True,
  212. )
  213. # add icon
  214. ToolTransformService.repack_provider(user_builtin_provider)
  215. tools = provider_controller.get_tools()
  216. for tool in tools:
  217. user_builtin_provider.tools.append(
  218. ToolTransformService.tool_to_user_tool(
  219. tenant_id=tenant_id,
  220. tool=tool,
  221. credentials=user_builtin_provider.original_credentials,
  222. labels=ToolLabelManager.get_tool_labels(provider_controller),
  223. )
  224. )
  225. result.append(user_builtin_provider)
  226. except Exception as e:
  227. raise e
  228. return BuiltinToolProviderSort.sort(result)