您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符

builtin_tools_manage_service.py 31KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750
  1. import json
  2. import logging
  3. import re
  4. from collections.abc import Mapping
  5. from pathlib import Path
  6. from typing import Any, Optional
  7. from sqlalchemy.orm import Session
  8. from configs import dify_config
  9. from constants import HIDDEN_VALUE, UNKNOWN_VALUE
  10. from core.helper.position_helper import is_filtered
  11. from core.helper.provider_cache import NoOpProviderCredentialCache, ToolProviderCredentialsCache
  12. from core.plugin.entities.plugin import ToolProviderID
  13. from core.tools.builtin_tool.provider import BuiltinToolProviderController
  14. from core.tools.builtin_tool.providers._positions import BuiltinToolProviderSort
  15. from core.tools.entities.api_entities import (
  16. ToolApiEntity,
  17. ToolProviderApiEntity,
  18. ToolProviderCredentialApiEntity,
  19. ToolProviderCredentialInfoApiEntity,
  20. )
  21. from core.tools.entities.tool_entities import CredentialType
  22. from core.tools.errors import ToolProviderNotFoundError
  23. from core.tools.plugin_tool.provider import PluginToolProviderController
  24. from core.tools.tool_label_manager import ToolLabelManager
  25. from core.tools.tool_manager import ToolManager
  26. from core.tools.utils.encryption import create_provider_encrypter
  27. from core.tools.utils.system_oauth_encryption import decrypt_system_oauth_params
  28. from extensions.ext_database import db
  29. from extensions.ext_redis import redis_client
  30. from models.tools import BuiltinToolProvider, ToolOAuthSystemClient, ToolOAuthTenantClient
  31. from services.plugin.plugin_service import PluginService
  32. from services.tools.tools_transform_service import ToolTransformService
  33. logger = logging.getLogger(__name__)
  34. class BuiltinToolManageService:
  35. __MAX_BUILTIN_TOOL_PROVIDER_COUNT__ = 100
  36. __DEFAULT_EXPIRES_AT__ = 2147483647
  37. @staticmethod
  38. def delete_custom_oauth_client_params(tenant_id: str, provider: str):
  39. """
  40. delete custom oauth client params
  41. """
  42. tool_provider = ToolProviderID(provider)
  43. with Session(db.engine) as session:
  44. session.query(ToolOAuthTenantClient).filter_by(
  45. tenant_id=tenant_id,
  46. provider=tool_provider.provider_name,
  47. plugin_id=tool_provider.plugin_id,
  48. ).delete()
  49. session.commit()
  50. return {"result": "success"}
  51. @staticmethod
  52. def get_builtin_tool_provider_oauth_client_schema(tenant_id: str, provider_name: str):
  53. """
  54. get builtin tool provider oauth client schema
  55. """
  56. provider = ToolManager.get_builtin_provider(provider_name, tenant_id)
  57. verified = not isinstance(provider, PluginToolProviderController) or PluginService.is_plugin_verified(
  58. tenant_id, provider.plugin_unique_identifier
  59. )
  60. is_oauth_custom_client_enabled = BuiltinToolManageService.is_oauth_custom_client_enabled(
  61. tenant_id, provider_name
  62. )
  63. is_system_oauth_params_exists = verified and BuiltinToolManageService.is_oauth_system_client_exists(
  64. provider_name
  65. )
  66. result = {
  67. "schema": provider.get_oauth_client_schema(),
  68. "is_oauth_custom_client_enabled": is_oauth_custom_client_enabled,
  69. "is_system_oauth_params_exists": is_system_oauth_params_exists,
  70. "client_params": BuiltinToolManageService.get_custom_oauth_client_params(tenant_id, provider_name),
  71. "redirect_uri": f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider_name}/tool/callback",
  72. }
  73. return result
  74. @staticmethod
  75. def list_builtin_tool_provider_tools(tenant_id: str, provider: str) -> list[ToolApiEntity]:
  76. """
  77. list builtin tool provider tools
  78. :param tenant_id: the id of the tenant
  79. :param provider: the name of the provider
  80. :return: the list of tools
  81. """
  82. provider_controller = ToolManager.get_builtin_provider(provider, tenant_id)
  83. tools = provider_controller.get_tools()
  84. result: list[ToolApiEntity] = []
  85. for tool in tools or []:
  86. result.append(
  87. ToolTransformService.convert_tool_entity_to_api_entity(
  88. tool=tool,
  89. tenant_id=tenant_id,
  90. labels=ToolLabelManager.get_tool_labels(provider_controller),
  91. )
  92. )
  93. return result
  94. @staticmethod
  95. def get_builtin_tool_provider_info(tenant_id: str, provider: str):
  96. """
  97. get builtin tool provider info
  98. """
  99. provider_controller = ToolManager.get_builtin_provider(provider, tenant_id)
  100. # check if user has added the provider
  101. builtin_provider = BuiltinToolManageService.get_builtin_provider(provider, tenant_id)
  102. if builtin_provider is None:
  103. raise ValueError(f"you have not added provider {provider}")
  104. entity = ToolTransformService.builtin_provider_to_user_provider(
  105. provider_controller=provider_controller,
  106. db_provider=builtin_provider,
  107. decrypt_credentials=True,
  108. )
  109. entity.original_credentials = {}
  110. return entity
  111. @staticmethod
  112. def list_builtin_provider_credentials_schema(provider_name: str, credential_type: CredentialType, tenant_id: str):
  113. """
  114. list builtin provider credentials schema
  115. :param credential_type: credential type
  116. :param provider_name: the name of the provider
  117. :param tenant_id: the id of the tenant
  118. :return: the list of tool providers
  119. """
  120. provider = ToolManager.get_builtin_provider(provider_name, tenant_id)
  121. return provider.get_credentials_schema_by_type(credential_type)
  122. @staticmethod
  123. def update_builtin_tool_provider(
  124. user_id: str,
  125. tenant_id: str,
  126. provider: str,
  127. credential_id: str,
  128. credentials: dict | None = None,
  129. name: str | None = None,
  130. ):
  131. """
  132. update builtin tool provider
  133. """
  134. with Session(db.engine) as session:
  135. # get if the provider exists
  136. db_provider = (
  137. session.query(BuiltinToolProvider)
  138. .where(
  139. BuiltinToolProvider.tenant_id == tenant_id,
  140. BuiltinToolProvider.id == credential_id,
  141. )
  142. .first()
  143. )
  144. if db_provider is None:
  145. raise ValueError(f"you have not added provider {provider}")
  146. try:
  147. if CredentialType.of(db_provider.credential_type).is_editable() and credentials:
  148. provider_controller = ToolManager.get_builtin_provider(provider, tenant_id)
  149. if not provider_controller.need_credentials:
  150. raise ValueError(f"provider {provider} does not need credentials")
  151. encrypter, cache = BuiltinToolManageService.create_tool_encrypter(
  152. tenant_id, db_provider, provider, provider_controller
  153. )
  154. original_credentials = encrypter.decrypt(db_provider.credentials)
  155. new_credentials: dict = {
  156. key: value if value != HIDDEN_VALUE else original_credentials.get(key, UNKNOWN_VALUE)
  157. for key, value in credentials.items()
  158. }
  159. if CredentialType.of(db_provider.credential_type).is_validate_allowed():
  160. provider_controller.validate_credentials(user_id, new_credentials)
  161. # encrypt credentials
  162. db_provider.encrypted_credentials = json.dumps(encrypter.encrypt(new_credentials))
  163. cache.delete()
  164. # update name if provided
  165. if name and name != db_provider.name:
  166. # check if the name is already used
  167. if (
  168. session.query(BuiltinToolProvider)
  169. .filter_by(tenant_id=tenant_id, provider=provider, name=name)
  170. .count()
  171. > 0
  172. ):
  173. raise ValueError(f"the credential name '{name}' is already used")
  174. db_provider.name = name
  175. session.commit()
  176. except Exception as e:
  177. session.rollback()
  178. raise ValueError(str(e))
  179. return {"result": "success"}
  180. @staticmethod
  181. def add_builtin_tool_provider(
  182. user_id: str,
  183. api_type: CredentialType,
  184. tenant_id: str,
  185. provider: str,
  186. credentials: dict,
  187. expires_at: int = -1,
  188. name: str | None = None,
  189. ):
  190. """
  191. add builtin tool provider
  192. """
  193. try:
  194. with Session(db.engine) as session:
  195. lock = f"builtin_tool_provider_create_lock:{tenant_id}_{provider}"
  196. with redis_client.lock(lock, timeout=20):
  197. provider_controller = ToolManager.get_builtin_provider(provider, tenant_id)
  198. if not provider_controller.need_credentials:
  199. raise ValueError(f"provider {provider} does not need credentials")
  200. provider_count = (
  201. session.query(BuiltinToolProvider).filter_by(tenant_id=tenant_id, provider=provider).count()
  202. )
  203. # check if the provider count is reached the limit
  204. if provider_count >= BuiltinToolManageService.__MAX_BUILTIN_TOOL_PROVIDER_COUNT__:
  205. raise ValueError(f"you have reached the maximum number of providers for {provider}")
  206. # validate credentials if allowed
  207. if CredentialType.of(api_type).is_validate_allowed():
  208. provider_controller.validate_credentials(user_id, credentials)
  209. # generate name if not provided
  210. if name is None or name == "":
  211. name = BuiltinToolManageService.generate_builtin_tool_provider_name(
  212. session=session, tenant_id=tenant_id, provider=provider, credential_type=api_type
  213. )
  214. else:
  215. # check if the name is already used
  216. if (
  217. session.query(BuiltinToolProvider)
  218. .filter_by(tenant_id=tenant_id, provider=provider, name=name)
  219. .count()
  220. > 0
  221. ):
  222. raise ValueError(f"the credential name '{name}' is already used")
  223. # create encrypter
  224. encrypter, _ = create_provider_encrypter(
  225. tenant_id=tenant_id,
  226. config=[
  227. x.to_basic_provider_config()
  228. for x in provider_controller.get_credentials_schema_by_type(api_type)
  229. ],
  230. cache=NoOpProviderCredentialCache(),
  231. )
  232. db_provider = BuiltinToolProvider(
  233. tenant_id=tenant_id,
  234. user_id=user_id,
  235. provider=provider,
  236. encrypted_credentials=json.dumps(encrypter.encrypt(credentials)),
  237. credential_type=api_type.value,
  238. name=name,
  239. expires_at=expires_at
  240. if expires_at is not None
  241. else BuiltinToolManageService.__DEFAULT_EXPIRES_AT__,
  242. )
  243. session.add(db_provider)
  244. session.commit()
  245. except Exception as e:
  246. session.rollback()
  247. raise ValueError(str(e))
  248. return {"result": "success"}
  249. @staticmethod
  250. def create_tool_encrypter(
  251. tenant_id: str,
  252. db_provider: BuiltinToolProvider,
  253. provider: str,
  254. provider_controller: BuiltinToolProviderController,
  255. ):
  256. encrypter, cache = create_provider_encrypter(
  257. tenant_id=tenant_id,
  258. config=[
  259. x.to_basic_provider_config()
  260. for x in provider_controller.get_credentials_schema_by_type(db_provider.credential_type)
  261. ],
  262. cache=ToolProviderCredentialsCache(tenant_id=tenant_id, provider=provider, credential_id=db_provider.id),
  263. )
  264. return encrypter, cache
  265. @staticmethod
  266. def generate_builtin_tool_provider_name(
  267. session: Session, tenant_id: str, provider: str, credential_type: CredentialType
  268. ) -> str:
  269. try:
  270. db_providers = (
  271. session.query(BuiltinToolProvider)
  272. .filter_by(
  273. tenant_id=tenant_id,
  274. provider=provider,
  275. credential_type=credential_type.value,
  276. )
  277. .order_by(BuiltinToolProvider.created_at.desc())
  278. .all()
  279. )
  280. # Get the default name pattern
  281. default_pattern = f"{credential_type.get_name()}"
  282. # Find all names that match the default pattern: "{default_pattern} {number}"
  283. pattern = rf"^{re.escape(default_pattern)}\s+(\d+)$"
  284. numbers = []
  285. for db_provider in db_providers:
  286. if db_provider.name:
  287. match = re.match(pattern, db_provider.name.strip())
  288. if match:
  289. numbers.append(int(match.group(1)))
  290. # If no default pattern names found, start with 1
  291. if not numbers:
  292. return f"{default_pattern} 1"
  293. # Find the next number
  294. max_number = max(numbers)
  295. return f"{default_pattern} {max_number + 1}"
  296. except Exception as e:
  297. logger.warning("Error generating next provider name for %s: %s", provider, str(e))
  298. # fallback
  299. return f"{credential_type.get_name()} 1"
  300. @staticmethod
  301. def get_builtin_tool_provider_credentials(
  302. tenant_id: str, provider_name: str
  303. ) -> list[ToolProviderCredentialApiEntity]:
  304. """
  305. get builtin tool provider credentials
  306. """
  307. with db.session.no_autoflush:
  308. providers = (
  309. db.session.query(BuiltinToolProvider)
  310. .filter_by(tenant_id=tenant_id, provider=provider_name)
  311. .order_by(BuiltinToolProvider.is_default.desc(), BuiltinToolProvider.created_at.asc())
  312. .all()
  313. )
  314. if len(providers) == 0:
  315. return []
  316. default_provider = providers[0]
  317. default_provider.is_default = True
  318. provider_controller = ToolManager.get_builtin_provider(default_provider.provider, tenant_id)
  319. credentials: list[ToolProviderCredentialApiEntity] = []
  320. encrypters = {}
  321. for provider in providers:
  322. credential_type = provider.credential_type
  323. if credential_type not in encrypters:
  324. encrypters[credential_type] = BuiltinToolManageService.create_tool_encrypter(
  325. tenant_id, provider, provider.provider, provider_controller
  326. )[0]
  327. encrypter = encrypters[credential_type]
  328. decrypt_credential = encrypter.mask_tool_credentials(encrypter.decrypt(provider.credentials))
  329. credential_entity = ToolTransformService.convert_builtin_provider_to_credential_entity(
  330. provider=provider,
  331. credentials=decrypt_credential,
  332. )
  333. credentials.append(credential_entity)
  334. return credentials
  335. @staticmethod
  336. def get_builtin_tool_provider_credential_info(tenant_id: str, provider: str) -> ToolProviderCredentialInfoApiEntity:
  337. """
  338. get builtin tool provider credential info
  339. """
  340. provider_controller = ToolManager.get_builtin_provider(provider, tenant_id)
  341. supported_credential_types = provider_controller.get_supported_credential_types()
  342. credentials = BuiltinToolManageService.get_builtin_tool_provider_credentials(tenant_id, provider)
  343. credential_info = ToolProviderCredentialInfoApiEntity(
  344. supported_credential_types=supported_credential_types,
  345. is_oauth_custom_client_enabled=BuiltinToolManageService.is_oauth_custom_client_enabled(tenant_id, provider),
  346. credentials=credentials,
  347. )
  348. return credential_info
  349. @staticmethod
  350. def delete_builtin_tool_provider(tenant_id: str, provider: str, credential_id: str):
  351. """
  352. delete tool provider
  353. """
  354. with Session(db.engine) as session:
  355. db_provider = (
  356. session.query(BuiltinToolProvider)
  357. .where(
  358. BuiltinToolProvider.tenant_id == tenant_id,
  359. BuiltinToolProvider.id == credential_id,
  360. )
  361. .first()
  362. )
  363. if db_provider is None:
  364. raise ValueError(f"you have not added provider {provider}")
  365. session.delete(db_provider)
  366. session.commit()
  367. # delete cache
  368. provider_controller = ToolManager.get_builtin_provider(provider, tenant_id)
  369. _, cache = BuiltinToolManageService.create_tool_encrypter(
  370. tenant_id, db_provider, provider, provider_controller
  371. )
  372. cache.delete()
  373. return {"result": "success"}
  374. @staticmethod
  375. def set_default_provider(tenant_id: str, user_id: str, provider: str, id: str):
  376. """
  377. set default provider
  378. """
  379. with Session(db.engine) as session:
  380. # get provider
  381. target_provider = session.query(BuiltinToolProvider).filter_by(id=id).first()
  382. if target_provider is None:
  383. raise ValueError("provider not found")
  384. # clear default provider
  385. session.query(BuiltinToolProvider).filter_by(
  386. tenant_id=tenant_id, user_id=user_id, provider=provider, is_default=True
  387. ).update({"is_default": False})
  388. # set new default provider
  389. target_provider.is_default = True
  390. session.commit()
  391. return {"result": "success"}
  392. @staticmethod
  393. def is_oauth_system_client_exists(provider_name: str) -> bool:
  394. """
  395. check if oauth system client exists
  396. """
  397. tool_provider = ToolProviderID(provider_name)
  398. with Session(db.engine).no_autoflush as session:
  399. system_client: ToolOAuthSystemClient | None = (
  400. session.query(ToolOAuthSystemClient)
  401. .filter_by(plugin_id=tool_provider.plugin_id, provider=tool_provider.provider_name)
  402. .first()
  403. )
  404. return system_client is not None
  405. @staticmethod
  406. def is_oauth_custom_client_enabled(tenant_id: str, provider: str) -> bool:
  407. """
  408. check if oauth custom client is enabled
  409. """
  410. tool_provider = ToolProviderID(provider)
  411. with Session(db.engine).no_autoflush as session:
  412. user_client: ToolOAuthTenantClient | None = (
  413. session.query(ToolOAuthTenantClient)
  414. .filter_by(
  415. tenant_id=tenant_id,
  416. provider=tool_provider.provider_name,
  417. plugin_id=tool_provider.plugin_id,
  418. enabled=True,
  419. )
  420. .first()
  421. )
  422. return user_client is not None and user_client.enabled
  423. @staticmethod
  424. def get_oauth_client(tenant_id: str, provider: str) -> Mapping[str, Any] | None:
  425. """
  426. get builtin tool provider
  427. """
  428. tool_provider = ToolProviderID(provider)
  429. provider_controller = ToolManager.get_builtin_provider(provider, tenant_id)
  430. encrypter, _ = create_provider_encrypter(
  431. tenant_id=tenant_id,
  432. config=[x.to_basic_provider_config() for x in provider_controller.get_oauth_client_schema()],
  433. cache=NoOpProviderCredentialCache(),
  434. )
  435. with Session(db.engine).no_autoflush as session:
  436. user_client: ToolOAuthTenantClient | None = (
  437. session.query(ToolOAuthTenantClient)
  438. .filter_by(
  439. tenant_id=tenant_id,
  440. provider=tool_provider.provider_name,
  441. plugin_id=tool_provider.plugin_id,
  442. enabled=True,
  443. )
  444. .first()
  445. )
  446. oauth_params: Mapping[str, Any] | None = None
  447. if user_client:
  448. oauth_params = encrypter.decrypt(user_client.oauth_params)
  449. return oauth_params
  450. # only verified provider can use official oauth client
  451. is_verified = not isinstance(
  452. provider_controller, PluginToolProviderController
  453. ) or PluginService.is_plugin_verified(tenant_id, provider_controller.plugin_unique_identifier)
  454. if not is_verified:
  455. return oauth_params
  456. system_client: ToolOAuthSystemClient | None = (
  457. session.query(ToolOAuthSystemClient)
  458. .filter_by(plugin_id=tool_provider.plugin_id, provider=tool_provider.provider_name)
  459. .first()
  460. )
  461. if system_client:
  462. try:
  463. oauth_params = decrypt_system_oauth_params(system_client.encrypted_oauth_params)
  464. except Exception as e:
  465. raise ValueError(f"Error decrypting system oauth params: {e}")
  466. return oauth_params
  467. @staticmethod
  468. def get_builtin_tool_provider_icon(provider: str):
  469. """
  470. get tool provider icon and it's mimetype
  471. """
  472. icon_path, mime_type = ToolManager.get_hardcoded_provider_icon(provider)
  473. icon_bytes = Path(icon_path).read_bytes()
  474. return icon_bytes, mime_type
  475. @staticmethod
  476. def list_builtin_tools(user_id: str, tenant_id: str) -> list[ToolProviderApiEntity]:
  477. """
  478. list builtin tools
  479. """
  480. # get all builtin providers
  481. provider_controllers = ToolManager.list_builtin_providers(tenant_id)
  482. with db.session.no_autoflush:
  483. # get all user added providers
  484. db_providers: list[BuiltinToolProvider] = ToolManager.list_default_builtin_providers(tenant_id)
  485. # rewrite db_providers
  486. for db_provider in db_providers:
  487. db_provider.provider = str(ToolProviderID(db_provider.provider))
  488. # find provider
  489. def find_provider(provider):
  490. return next(filter(lambda db_provider: db_provider.provider == provider, db_providers), None)
  491. result: list[ToolProviderApiEntity] = []
  492. for provider_controller in provider_controllers:
  493. try:
  494. # handle include, exclude
  495. if is_filtered(
  496. include_set=dify_config.POSITION_TOOL_INCLUDES_SET, # type: ignore
  497. exclude_set=dify_config.POSITION_TOOL_EXCLUDES_SET, # type: ignore
  498. data=provider_controller,
  499. name_func=lambda x: x.identity.name,
  500. ):
  501. continue
  502. # convert provider controller to user provider
  503. user_builtin_provider = ToolTransformService.builtin_provider_to_user_provider(
  504. provider_controller=provider_controller,
  505. db_provider=find_provider(provider_controller.entity.identity.name),
  506. decrypt_credentials=True,
  507. )
  508. # add icon
  509. ToolTransformService.repack_provider(tenant_id=tenant_id, provider=user_builtin_provider)
  510. tools = provider_controller.get_tools()
  511. for tool in tools or []:
  512. user_builtin_provider.tools.append(
  513. ToolTransformService.convert_tool_entity_to_api_entity(
  514. tenant_id=tenant_id,
  515. tool=tool,
  516. labels=ToolLabelManager.get_tool_labels(provider_controller),
  517. )
  518. )
  519. result.append(user_builtin_provider)
  520. except Exception as e:
  521. raise e
  522. return BuiltinToolProviderSort.sort(result)
  523. @staticmethod
  524. def get_builtin_provider(provider_name: str, tenant_id: str) -> Optional[BuiltinToolProvider]:
  525. """
  526. This method is used to fetch the builtin provider from the database
  527. 1.if the default provider exists, return the default provider
  528. 2.if the default provider does not exist, return the oldest provider
  529. """
  530. with Session(db.engine) as session:
  531. try:
  532. full_provider_name = provider_name
  533. provider_id_entity = ToolProviderID(provider_name)
  534. provider_name = provider_id_entity.provider_name
  535. if provider_id_entity.organization != "langgenius":
  536. provider = (
  537. session.query(BuiltinToolProvider)
  538. .where(
  539. BuiltinToolProvider.tenant_id == tenant_id,
  540. BuiltinToolProvider.provider == full_provider_name,
  541. )
  542. .order_by(
  543. BuiltinToolProvider.is_default.desc(), # default=True first
  544. BuiltinToolProvider.created_at.asc(), # oldest first
  545. )
  546. .first()
  547. )
  548. else:
  549. provider = (
  550. session.query(BuiltinToolProvider)
  551. .where(
  552. BuiltinToolProvider.tenant_id == tenant_id,
  553. (BuiltinToolProvider.provider == provider_name)
  554. | (BuiltinToolProvider.provider == full_provider_name),
  555. )
  556. .order_by(
  557. BuiltinToolProvider.is_default.desc(), # default=True first
  558. BuiltinToolProvider.created_at.asc(), # oldest first
  559. )
  560. .first()
  561. )
  562. if provider is None:
  563. return None
  564. provider.provider = ToolProviderID(provider.provider).to_string()
  565. return provider
  566. except Exception:
  567. # it's an old provider without organization
  568. return (
  569. session.query(BuiltinToolProvider)
  570. .where(BuiltinToolProvider.tenant_id == tenant_id, BuiltinToolProvider.provider == provider_name)
  571. .order_by(
  572. BuiltinToolProvider.is_default.desc(), # default=True first
  573. BuiltinToolProvider.created_at.asc(), # oldest first
  574. )
  575. .first()
  576. )
  577. @staticmethod
  578. def save_custom_oauth_client_params(
  579. tenant_id: str,
  580. provider: str,
  581. client_params: Optional[dict] = None,
  582. enable_oauth_custom_client: Optional[bool] = None,
  583. ):
  584. """
  585. setup oauth custom client
  586. """
  587. if client_params is None and enable_oauth_custom_client is None:
  588. return {"result": "success"}
  589. tool_provider = ToolProviderID(provider)
  590. provider_controller = ToolManager.get_builtin_provider(provider, tenant_id)
  591. if not provider_controller:
  592. raise ToolProviderNotFoundError(f"Provider {provider} not found")
  593. if not isinstance(provider_controller, (BuiltinToolProviderController, PluginToolProviderController)):
  594. raise ValueError(f"Provider {provider} is not a builtin or plugin provider")
  595. with Session(db.engine) as session:
  596. custom_client_params = (
  597. session.query(ToolOAuthTenantClient)
  598. .filter_by(
  599. tenant_id=tenant_id,
  600. plugin_id=tool_provider.plugin_id,
  601. provider=tool_provider.provider_name,
  602. )
  603. .first()
  604. )
  605. # if the record does not exist, create a basic record
  606. if custom_client_params is None:
  607. custom_client_params = ToolOAuthTenantClient(
  608. tenant_id=tenant_id,
  609. plugin_id=tool_provider.plugin_id,
  610. provider=tool_provider.provider_name,
  611. )
  612. session.add(custom_client_params)
  613. if client_params is not None:
  614. encrypter, _ = create_provider_encrypter(
  615. tenant_id=tenant_id,
  616. config=[x.to_basic_provider_config() for x in provider_controller.get_oauth_client_schema()],
  617. cache=NoOpProviderCredentialCache(),
  618. )
  619. original_params = encrypter.decrypt(custom_client_params.oauth_params)
  620. new_params: dict = {
  621. key: value if value != HIDDEN_VALUE else original_params.get(key, UNKNOWN_VALUE)
  622. for key, value in client_params.items()
  623. }
  624. custom_client_params.encrypted_oauth_params = json.dumps(encrypter.encrypt(new_params))
  625. if enable_oauth_custom_client is not None:
  626. custom_client_params.enabled = enable_oauth_custom_client
  627. session.commit()
  628. return {"result": "success"}
  629. @staticmethod
  630. def get_custom_oauth_client_params(tenant_id: str, provider: str):
  631. """
  632. get custom oauth client params
  633. """
  634. with Session(db.engine) as session:
  635. tool_provider = ToolProviderID(provider)
  636. custom_oauth_client_params: ToolOAuthTenantClient | None = (
  637. session.query(ToolOAuthTenantClient)
  638. .filter_by(
  639. tenant_id=tenant_id,
  640. plugin_id=tool_provider.plugin_id,
  641. provider=tool_provider.provider_name,
  642. )
  643. .first()
  644. )
  645. if custom_oauth_client_params is None:
  646. return {}
  647. provider_controller = ToolManager.get_builtin_provider(provider, tenant_id)
  648. if not provider_controller:
  649. raise ToolProviderNotFoundError(f"Provider {provider} not found")
  650. if not isinstance(provider_controller, BuiltinToolProviderController):
  651. raise ValueError(f"Provider {provider} is not a builtin or plugin provider")
  652. encrypter, _ = create_provider_encrypter(
  653. tenant_id=tenant_id,
  654. config=[x.to_basic_provider_config() for x in provider_controller.get_oauth_client_schema()],
  655. cache=NoOpProviderCredentialCache(),
  656. )
  657. return encrypter.mask_tool_credentials(encrypter.decrypt(custom_oauth_client_params.oauth_params))