provider_manager.py 40 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926
  1. import json
  2. from collections import defaultdict
  3. from json import JSONDecodeError
  4. from typing import Optional
  5. from sqlalchemy.exc import IntegrityError
  6. from configs import dify_config
  7. from core.entities.model_entities import DefaultModelEntity, DefaultModelProviderEntity
  8. from core.entities.provider_configuration import ProviderConfiguration, ProviderConfigurations, ProviderModelBundle
  9. from core.entities.provider_entities import (
  10. CustomConfiguration,
  11. CustomModelConfiguration,
  12. CustomProviderConfiguration,
  13. ModelLoadBalancingConfiguration,
  14. ModelSettings,
  15. ProviderQuotaType,
  16. QuotaConfiguration,
  17. SystemConfiguration,
  18. )
  19. from core.helper import encrypter
  20. from core.helper.model_provider_cache import ProviderCredentialsCache, ProviderCredentialsCacheType
  21. from core.helper.position_helper import is_filtered
  22. from core.model_runtime.entities.model_entities import ModelType
  23. from core.model_runtime.entities.provider_entities import CredentialFormSchema, FormType, ProviderEntity
  24. from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
  25. from extensions import ext_hosting_provider
  26. from extensions.ext_database import db
  27. from extensions.ext_redis import redis_client
  28. from models.provider import (
  29. LoadBalancingModelConfig,
  30. Provider,
  31. ProviderModel,
  32. ProviderModelSetting,
  33. ProviderType,
  34. TenantDefaultModel,
  35. TenantPreferredModelProvider,
  36. )
  37. from services.feature_service import FeatureService
  38. class ProviderManager:
  39. """
  40. ProviderManager is a class that manages the model providers includes Hosting and Customize Model Providers.
  41. """
  42. def __init__(self) -> None:
  43. self.decoding_rsa_key = None
  44. self.decoding_cipher_rsa = None
  45. def get_configurations(self, tenant_id: str) -> ProviderConfigurations:
  46. """
  47. Get model provider configurations.
  48. Construct ProviderConfiguration objects for each provider
  49. Including:
  50. 1. Basic information of the provider
  51. 2. Hosting configuration information, including:
  52. (1. Whether to enable (support) hosting type, if enabled, the following information exists
  53. (2. List of hosting type provider configurations
  54. (including quota type, quota limit, current remaining quota, etc.)
  55. (3. The current hosting type in use (whether there is a quota or not)
  56. paid quotas > provider free quotas > hosting trial quotas
  57. (4. Unified credentials for hosting providers
  58. 3. Custom configuration information, including:
  59. (1. Whether to enable (support) custom type, if enabled, the following information exists
  60. (2. Custom provider configuration (including credentials)
  61. (3. List of custom provider model configurations (including credentials)
  62. 4. Hosting/custom preferred provider type.
  63. Provide methods:
  64. - Get the current configuration (including credentials)
  65. - Get the availability and status of the hosting configuration: active available,
  66. quota_exceeded insufficient quota, unsupported hosting
  67. - Get the availability of custom configuration
  68. Custom provider available conditions:
  69. (1. custom provider credentials available
  70. (2. at least one custom model credentials available
  71. - Verify, update, and delete custom provider configuration
  72. - Verify, update, and delete custom provider model configuration
  73. - Get the list of available models (optional provider filtering, model type filtering)
  74. Append custom provider models to the list
  75. - Get provider instance
  76. - Switch selection priority
  77. :param tenant_id:
  78. :return:
  79. """
  80. # Get all provider records of the workspace
  81. provider_name_to_provider_records_dict = self._get_all_providers(tenant_id)
  82. # Initialize trial provider records if not exist
  83. provider_name_to_provider_records_dict = self._init_trial_provider_records(
  84. tenant_id, provider_name_to_provider_records_dict
  85. )
  86. # Get all provider model records of the workspace
  87. provider_name_to_provider_model_records_dict = self._get_all_provider_models(tenant_id)
  88. # Get all provider entities
  89. model_provider_factory = ModelProviderFactory(tenant_id)
  90. provider_entities = model_provider_factory.get_providers()
  91. # Get All preferred provider types of the workspace
  92. provider_name_to_preferred_model_provider_records_dict = self._get_all_preferred_model_providers(tenant_id)
  93. # Get All provider model settings
  94. provider_name_to_provider_model_settings_dict = self._get_all_provider_model_settings(tenant_id)
  95. # Get All load balancing configs
  96. provider_name_to_provider_load_balancing_model_configs_dict = self._get_all_provider_load_balancing_configs(
  97. tenant_id
  98. )
  99. provider_configurations = ProviderConfigurations(tenant_id=tenant_id)
  100. # Construct ProviderConfiguration objects for each provider
  101. for provider_entity in provider_entities:
  102. # handle include, exclude
  103. if is_filtered(
  104. include_set=dify_config.POSITION_PROVIDER_INCLUDES_SET,
  105. exclude_set=dify_config.POSITION_PROVIDER_EXCLUDES_SET,
  106. data=provider_entity,
  107. name_func=lambda x: x.provider,
  108. ):
  109. continue
  110. provider_name = provider_entity.provider
  111. provider_records = provider_name_to_provider_records_dict.get(provider_entity.provider, [])
  112. provider_model_records = provider_name_to_provider_model_records_dict.get(provider_entity.provider, [])
  113. # Convert to custom configuration
  114. custom_configuration = self._to_custom_configuration(
  115. tenant_id, provider_entity, provider_records, provider_model_records
  116. )
  117. # Convert to system configuration
  118. system_configuration = self._to_system_configuration(tenant_id, provider_entity, provider_records)
  119. # Get preferred provider type
  120. preferred_provider_type_record = provider_name_to_preferred_model_provider_records_dict.get(provider_name)
  121. if preferred_provider_type_record:
  122. preferred_provider_type = ProviderType.value_of(preferred_provider_type_record.preferred_provider_type)
  123. elif custom_configuration.provider or custom_configuration.models:
  124. preferred_provider_type = ProviderType.CUSTOM
  125. elif system_configuration.enabled:
  126. preferred_provider_type = ProviderType.SYSTEM
  127. else:
  128. preferred_provider_type = ProviderType.CUSTOM
  129. using_provider_type = preferred_provider_type
  130. has_valid_quota = any(quota_conf.is_valid for quota_conf in system_configuration.quota_configurations)
  131. if preferred_provider_type == ProviderType.SYSTEM:
  132. if not system_configuration.enabled or not has_valid_quota:
  133. using_provider_type = ProviderType.CUSTOM
  134. else:
  135. if not custom_configuration.provider and not custom_configuration.models:
  136. if system_configuration.enabled and has_valid_quota:
  137. using_provider_type = ProviderType.SYSTEM
  138. # Get provider load balancing configs
  139. provider_model_settings = provider_name_to_provider_model_settings_dict.get(provider_name)
  140. # Get provider load balancing configs
  141. provider_load_balancing_configs = provider_name_to_provider_load_balancing_model_configs_dict.get(
  142. provider_name
  143. )
  144. # Convert to model settings
  145. model_settings = self._to_model_settings(
  146. provider_entity=provider_entity,
  147. provider_model_settings=provider_model_settings,
  148. load_balancing_model_configs=provider_load_balancing_configs,
  149. )
  150. provider_configuration = ProviderConfiguration(
  151. tenant_id=tenant_id,
  152. provider=provider_entity,
  153. preferred_provider_type=preferred_provider_type,
  154. using_provider_type=using_provider_type,
  155. system_configuration=system_configuration,
  156. custom_configuration=custom_configuration,
  157. model_settings=model_settings,
  158. )
  159. provider_configurations[provider_name] = provider_configuration
  160. # Return the encapsulated object
  161. return provider_configurations
  162. def get_provider_model_bundle(self, tenant_id: str, provider: str, model_type: ModelType) -> ProviderModelBundle:
  163. """
  164. Get provider model bundle.
  165. :param tenant_id: workspace id
  166. :param provider: provider name
  167. :param model_type: model type
  168. :return:
  169. """
  170. provider_configurations = self.get_configurations(tenant_id)
  171. # get provider instance
  172. provider_configuration = provider_configurations.get(provider)
  173. if not provider_configuration:
  174. raise ValueError(f"Provider {provider} does not exist.")
  175. model_type_instance = provider_configuration.get_model_type_instance(model_type)
  176. return ProviderModelBundle(
  177. configuration=provider_configuration,
  178. model_type_instance=model_type_instance,
  179. )
  180. def get_default_model(self, tenant_id: str, model_type: ModelType) -> Optional[DefaultModelEntity]:
  181. """
  182. Get default model.
  183. :param tenant_id: workspace id
  184. :param model_type: model type
  185. :return:
  186. """
  187. # Get the corresponding TenantDefaultModel record
  188. default_model = (
  189. db.session.query(TenantDefaultModel)
  190. .filter(
  191. TenantDefaultModel.tenant_id == tenant_id,
  192. TenantDefaultModel.model_type == model_type.to_origin_model_type(),
  193. )
  194. .first()
  195. )
  196. # If it does not exist, get the first available provider model from get_configurations
  197. # and update the TenantDefaultModel record
  198. if not default_model:
  199. # Get provider configurations
  200. provider_configurations = self.get_configurations(tenant_id)
  201. # get available models from provider_configurations
  202. available_models = provider_configurations.get_models(model_type=model_type, only_active=True)
  203. if available_models:
  204. available_model = next(
  205. (model for model in available_models if model.model == "gpt-4"), available_models[0]
  206. )
  207. default_model = TenantDefaultModel()
  208. default_model.tenant_id = tenant_id
  209. default_model.model_type = model_type.to_origin_model_type()
  210. default_model.provider_name = available_model.provider.provider
  211. default_model.model_name = available_model.model
  212. db.session.add(default_model)
  213. db.session.commit()
  214. if not default_model:
  215. return None
  216. model_provider_factory = ModelProviderFactory(tenant_id)
  217. provider_schema = model_provider_factory.get_provider_schema(provider=default_model.provider_name)
  218. return DefaultModelEntity(
  219. model=default_model.model_name,
  220. model_type=model_type,
  221. provider=DefaultModelProviderEntity(
  222. provider=provider_schema.provider,
  223. label=provider_schema.label,
  224. icon_small=provider_schema.icon_small,
  225. icon_large=provider_schema.icon_large,
  226. supported_model_types=provider_schema.supported_model_types,
  227. ),
  228. )
  229. def get_first_provider_first_model(self, tenant_id: str, model_type: ModelType) -> tuple[str | None, str | None]:
  230. """
  231. Get names of first model and its provider
  232. :param tenant_id: workspace id
  233. :param model_type: model type
  234. :return: provider name, model name
  235. """
  236. provider_configurations = self.get_configurations(tenant_id)
  237. # get available models from provider_configurations
  238. all_models = provider_configurations.get_models(model_type=model_type, only_active=False)
  239. if not all_models:
  240. return None, None
  241. return all_models[0].provider.provider, all_models[0].model
  242. def update_default_model_record(
  243. self, tenant_id: str, model_type: ModelType, provider: str, model: str
  244. ) -> TenantDefaultModel:
  245. """
  246. Update default model record.
  247. :param tenant_id: workspace id
  248. :param model_type: model type
  249. :param provider: provider name
  250. :param model: model name
  251. :return:
  252. """
  253. provider_configurations = self.get_configurations(tenant_id)
  254. if provider not in provider_configurations:
  255. raise ValueError(f"Provider {provider} does not exist.")
  256. # get available models from provider_configurations
  257. available_models = provider_configurations.get_models(model_type=model_type, only_active=True)
  258. # check if the model is exist in available models
  259. model_names = [model.model for model in available_models]
  260. if model not in model_names:
  261. raise ValueError(f"Model {model} does not exist.")
  262. # Get the list of available models from get_configurations and check if it is LLM
  263. default_model = (
  264. db.session.query(TenantDefaultModel)
  265. .filter(
  266. TenantDefaultModel.tenant_id == tenant_id,
  267. TenantDefaultModel.model_type == model_type.to_origin_model_type(),
  268. )
  269. .first()
  270. )
  271. # create or update TenantDefaultModel record
  272. if default_model:
  273. # update default model
  274. default_model.provider_name = provider
  275. default_model.model_name = model
  276. db.session.commit()
  277. else:
  278. # create default model
  279. default_model = TenantDefaultModel(
  280. tenant_id=tenant_id,
  281. model_type=model_type.value,
  282. provider_name=provider,
  283. model_name=model,
  284. )
  285. db.session.add(default_model)
  286. db.session.commit()
  287. return default_model
  288. @staticmethod
  289. def _get_all_providers(tenant_id: str) -> dict[str, list[Provider]]:
  290. """
  291. Get all provider records of the workspace.
  292. :param tenant_id: workspace id
  293. :return:
  294. """
  295. providers = db.session.query(Provider).filter(Provider.tenant_id == tenant_id, Provider.is_valid == True).all()
  296. provider_name_to_provider_records_dict = defaultdict(list)
  297. for provider in providers:
  298. provider_name_to_provider_records_dict[provider.provider_name].append(provider)
  299. return provider_name_to_provider_records_dict
  300. @staticmethod
  301. def _get_all_provider_models(tenant_id: str) -> dict[str, list[ProviderModel]]:
  302. """
  303. Get all provider model records of the workspace.
  304. :param tenant_id: workspace id
  305. :return:
  306. """
  307. # Get all provider model records of the workspace
  308. provider_models = (
  309. db.session.query(ProviderModel)
  310. .filter(ProviderModel.tenant_id == tenant_id, ProviderModel.is_valid == True)
  311. .all()
  312. )
  313. provider_name_to_provider_model_records_dict = defaultdict(list)
  314. for provider_model in provider_models:
  315. provider_name_to_provider_model_records_dict[provider_model.provider_name].append(provider_model)
  316. return provider_name_to_provider_model_records_dict
  317. @staticmethod
  318. def _get_all_preferred_model_providers(tenant_id: str) -> dict[str, TenantPreferredModelProvider]:
  319. """
  320. Get All preferred provider types of the workspace.
  321. :param tenant_id: workspace id
  322. :return:
  323. """
  324. preferred_provider_types = (
  325. db.session.query(TenantPreferredModelProvider)
  326. .filter(TenantPreferredModelProvider.tenant_id == tenant_id)
  327. .all()
  328. )
  329. provider_name_to_preferred_provider_type_records_dict = {
  330. preferred_provider_type.provider_name: preferred_provider_type
  331. for preferred_provider_type in preferred_provider_types
  332. }
  333. return provider_name_to_preferred_provider_type_records_dict
  334. @staticmethod
  335. def _get_all_provider_model_settings(tenant_id: str) -> dict[str, list[ProviderModelSetting]]:
  336. """
  337. Get All provider model settings of the workspace.
  338. :param tenant_id: workspace id
  339. :return:
  340. """
  341. provider_model_settings = (
  342. db.session.query(ProviderModelSetting).filter(ProviderModelSetting.tenant_id == tenant_id).all()
  343. )
  344. provider_name_to_provider_model_settings_dict = defaultdict(list)
  345. for provider_model_setting in provider_model_settings:
  346. (
  347. provider_name_to_provider_model_settings_dict[provider_model_setting.provider_name].append(
  348. provider_model_setting
  349. )
  350. )
  351. return provider_name_to_provider_model_settings_dict
  352. @staticmethod
  353. def _get_all_provider_load_balancing_configs(tenant_id: str) -> dict[str, list[LoadBalancingModelConfig]]:
  354. """
  355. Get All provider load balancing configs of the workspace.
  356. :param tenant_id: workspace id
  357. :return:
  358. """
  359. cache_key = f"tenant:{tenant_id}:model_load_balancing_enabled"
  360. cache_result = redis_client.get(cache_key)
  361. if cache_result is None:
  362. model_load_balancing_enabled = FeatureService.get_features(tenant_id).model_load_balancing_enabled
  363. redis_client.setex(cache_key, 120, str(model_load_balancing_enabled))
  364. else:
  365. cache_result = cache_result.decode("utf-8")
  366. model_load_balancing_enabled = cache_result == "True"
  367. if not model_load_balancing_enabled:
  368. return {}
  369. provider_load_balancing_configs = (
  370. db.session.query(LoadBalancingModelConfig).filter(LoadBalancingModelConfig.tenant_id == tenant_id).all()
  371. )
  372. provider_name_to_provider_load_balancing_model_configs_dict = defaultdict(list)
  373. for provider_load_balancing_config in provider_load_balancing_configs:
  374. (
  375. provider_name_to_provider_load_balancing_model_configs_dict[
  376. provider_load_balancing_config.provider_name
  377. ].append(provider_load_balancing_config)
  378. )
  379. return provider_name_to_provider_load_balancing_model_configs_dict
  380. @staticmethod
  381. def _init_trial_provider_records(
  382. tenant_id: str, provider_name_to_provider_records_dict: dict[str, list]
  383. ) -> dict[str, list]:
  384. """
  385. Initialize trial provider records if not exists.
  386. :param tenant_id: workspace id
  387. :param provider_name_to_provider_records_dict: provider name to provider records dict
  388. :return:
  389. """
  390. # Get hosting configuration
  391. hosting_configuration = ext_hosting_provider.hosting_configuration
  392. for provider_name, configuration in hosting_configuration.provider_map.items():
  393. if not configuration.enabled:
  394. continue
  395. provider_records = provider_name_to_provider_records_dict.get(provider_name)
  396. if not provider_records:
  397. provider_records = []
  398. provider_quota_to_provider_record_dict = {}
  399. for provider_record in provider_records:
  400. if provider_record.provider_type != ProviderType.SYSTEM.value:
  401. continue
  402. provider_quota_to_provider_record_dict[ProviderQuotaType.value_of(provider_record.quota_type)] = (
  403. provider_record
  404. )
  405. for quota in configuration.quotas:
  406. if quota.quota_type == ProviderQuotaType.TRIAL:
  407. # Init trial provider records if not exists
  408. if ProviderQuotaType.TRIAL not in provider_quota_to_provider_record_dict:
  409. try:
  410. provider_record = Provider()
  411. provider_record.tenant_id = tenant_id
  412. provider_record.provider_name = provider_name
  413. provider_record.provider_type = ProviderType.SYSTEM.value
  414. provider_record.quota_type = ProviderQuotaType.TRIAL.value
  415. provider_record.quota_limit = quota.quota_limit
  416. provider_record.quota_used = 0
  417. provider_record.is_valid = True
  418. db.session.add(provider_record)
  419. db.session.commit()
  420. except IntegrityError:
  421. db.session.rollback()
  422. provider_record = (
  423. db.session.query(Provider)
  424. .filter(
  425. Provider.tenant_id == tenant_id,
  426. Provider.provider_name == provider_name,
  427. Provider.provider_type == ProviderType.SYSTEM.value,
  428. Provider.quota_type == ProviderQuotaType.TRIAL.value,
  429. )
  430. .first()
  431. )
  432. if provider_record and not provider_record.is_valid:
  433. provider_record.is_valid = True
  434. db.session.commit()
  435. provider_name_to_provider_records_dict[provider_name].append(provider_record)
  436. return provider_name_to_provider_records_dict
  437. def _to_custom_configuration(
  438. self,
  439. tenant_id: str,
  440. provider_entity: ProviderEntity,
  441. provider_records: list[Provider],
  442. provider_model_records: list[ProviderModel],
  443. ) -> CustomConfiguration:
  444. """
  445. Convert to custom configuration.
  446. :param tenant_id: workspace id
  447. :param provider_entity: provider entity
  448. :param provider_records: provider records
  449. :param provider_model_records: provider model records
  450. :return:
  451. """
  452. # Get provider credential secret variables
  453. provider_credential_secret_variables = self._extract_secret_variables(
  454. provider_entity.provider_credential_schema.credential_form_schemas
  455. if provider_entity.provider_credential_schema
  456. else []
  457. )
  458. # Get custom provider record
  459. custom_provider_record = None
  460. for provider_record in provider_records:
  461. if provider_record.provider_type == ProviderType.SYSTEM.value:
  462. continue
  463. if not provider_record.encrypted_config:
  464. continue
  465. custom_provider_record = provider_record
  466. # Get custom provider credentials
  467. custom_provider_configuration = None
  468. if custom_provider_record:
  469. provider_credentials_cache = ProviderCredentialsCache(
  470. tenant_id=tenant_id,
  471. identity_id=custom_provider_record.id,
  472. cache_type=ProviderCredentialsCacheType.PROVIDER,
  473. )
  474. # Get cached provider credentials
  475. cached_provider_credentials = provider_credentials_cache.get()
  476. if not cached_provider_credentials:
  477. try:
  478. # fix origin data
  479. if (
  480. custom_provider_record.encrypted_config
  481. and not custom_provider_record.encrypted_config.startswith("{")
  482. ):
  483. provider_credentials = {"openai_api_key": custom_provider_record.encrypted_config}
  484. else:
  485. provider_credentials = json.loads(custom_provider_record.encrypted_config)
  486. except JSONDecodeError:
  487. provider_credentials = {}
  488. # Get decoding rsa key and cipher for decrypting credentials
  489. if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None:
  490. self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id)
  491. for variable in provider_credential_secret_variables:
  492. if variable in provider_credentials:
  493. try:
  494. provider_credentials[variable] = encrypter.decrypt_token_with_decoding(
  495. provider_credentials.get(variable), self.decoding_rsa_key, self.decoding_cipher_rsa
  496. )
  497. except ValueError:
  498. pass
  499. # cache provider credentials
  500. provider_credentials_cache.set(credentials=provider_credentials)
  501. else:
  502. provider_credentials = cached_provider_credentials
  503. custom_provider_configuration = CustomProviderConfiguration(credentials=provider_credentials)
  504. # Get provider model credential secret variables
  505. model_credential_secret_variables = self._extract_secret_variables(
  506. provider_entity.model_credential_schema.credential_form_schemas
  507. if provider_entity.model_credential_schema
  508. else []
  509. )
  510. # Get custom provider model credentials
  511. custom_model_configurations = []
  512. for provider_model_record in provider_model_records:
  513. if not provider_model_record.encrypted_config:
  514. continue
  515. provider_model_credentials_cache = ProviderCredentialsCache(
  516. tenant_id=tenant_id, identity_id=provider_model_record.id, cache_type=ProviderCredentialsCacheType.MODEL
  517. )
  518. # Get cached provider model credentials
  519. cached_provider_model_credentials = provider_model_credentials_cache.get()
  520. if not cached_provider_model_credentials:
  521. try:
  522. provider_model_credentials = json.loads(provider_model_record.encrypted_config)
  523. except JSONDecodeError:
  524. continue
  525. # Get decoding rsa key and cipher for decrypting credentials
  526. if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None:
  527. self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id)
  528. for variable in model_credential_secret_variables:
  529. if variable in provider_model_credentials:
  530. try:
  531. provider_model_credentials[variable] = encrypter.decrypt_token_with_decoding(
  532. provider_model_credentials.get(variable),
  533. self.decoding_rsa_key,
  534. self.decoding_cipher_rsa,
  535. )
  536. except ValueError:
  537. pass
  538. # cache provider model credentials
  539. provider_model_credentials_cache.set(credentials=provider_model_credentials)
  540. else:
  541. provider_model_credentials = cached_provider_model_credentials
  542. custom_model_configurations.append(
  543. CustomModelConfiguration(
  544. model=provider_model_record.model_name,
  545. model_type=ModelType.value_of(provider_model_record.model_type),
  546. credentials=provider_model_credentials,
  547. )
  548. )
  549. return CustomConfiguration(provider=custom_provider_configuration, models=custom_model_configurations)
  550. def _to_system_configuration(
  551. self, tenant_id: str, provider_entity: ProviderEntity, provider_records: list[Provider]
  552. ) -> SystemConfiguration:
  553. """
  554. Convert to system configuration.
  555. :param tenant_id: workspace id
  556. :param provider_entity: provider entity
  557. :param provider_records: provider records
  558. :return:
  559. """
  560. # Get hosting configuration
  561. hosting_configuration = ext_hosting_provider.hosting_configuration
  562. if (
  563. provider_entity.provider not in hosting_configuration.provider_map
  564. or not hosting_configuration.provider_map.get(provider_entity.provider).enabled
  565. ):
  566. return SystemConfiguration(enabled=False)
  567. provider_hosting_configuration = hosting_configuration.provider_map.get(provider_entity.provider)
  568. # Convert provider_records to dict
  569. quota_type_to_provider_records_dict = {}
  570. for provider_record in provider_records:
  571. if provider_record.provider_type != ProviderType.SYSTEM.value:
  572. continue
  573. quota_type_to_provider_records_dict[ProviderQuotaType.value_of(provider_record.quota_type)] = (
  574. provider_record
  575. )
  576. quota_configurations = []
  577. for provider_quota in provider_hosting_configuration.quotas:
  578. if provider_quota.quota_type not in quota_type_to_provider_records_dict:
  579. if provider_quota.quota_type == ProviderQuotaType.FREE:
  580. quota_configuration = QuotaConfiguration(
  581. quota_type=provider_quota.quota_type,
  582. quota_unit=provider_hosting_configuration.quota_unit,
  583. quota_used=0,
  584. quota_limit=0,
  585. is_valid=False,
  586. restrict_models=provider_quota.restrict_models,
  587. )
  588. else:
  589. continue
  590. else:
  591. provider_record = quota_type_to_provider_records_dict[provider_quota.quota_type]
  592. quota_configuration = QuotaConfiguration(
  593. quota_type=provider_quota.quota_type,
  594. quota_unit=provider_hosting_configuration.quota_unit,
  595. quota_used=provider_record.quota_used,
  596. quota_limit=provider_record.quota_limit,
  597. is_valid=provider_record.quota_limit > provider_record.quota_used
  598. or provider_record.quota_limit == -1,
  599. restrict_models=provider_quota.restrict_models,
  600. )
  601. quota_configurations.append(quota_configuration)
  602. if len(quota_configurations) == 0:
  603. return SystemConfiguration(enabled=False)
  604. current_quota_type = self._choice_current_using_quota_type(quota_configurations)
  605. current_using_credentials = provider_hosting_configuration.credentials
  606. if current_quota_type == ProviderQuotaType.FREE:
  607. provider_record = quota_type_to_provider_records_dict.get(current_quota_type)
  608. if provider_record:
  609. provider_credentials_cache = ProviderCredentialsCache(
  610. tenant_id=tenant_id,
  611. identity_id=provider_record.id,
  612. cache_type=ProviderCredentialsCacheType.PROVIDER,
  613. )
  614. # Get cached provider credentials
  615. cached_provider_credentials = provider_credentials_cache.get()
  616. if not cached_provider_credentials:
  617. try:
  618. provider_credentials = json.loads(provider_record.encrypted_config)
  619. except JSONDecodeError:
  620. provider_credentials = {}
  621. # Get provider credential secret variables
  622. provider_credential_secret_variables = self._extract_secret_variables(
  623. provider_entity.provider_credential_schema.credential_form_schemas
  624. if provider_entity.provider_credential_schema
  625. else []
  626. )
  627. # Get decoding rsa key and cipher for decrypting credentials
  628. if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None:
  629. self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id)
  630. for variable in provider_credential_secret_variables:
  631. if variable in provider_credentials:
  632. try:
  633. provider_credentials[variable] = encrypter.decrypt_token_with_decoding(
  634. provider_credentials.get(variable), self.decoding_rsa_key, self.decoding_cipher_rsa
  635. )
  636. except ValueError:
  637. pass
  638. current_using_credentials = provider_credentials
  639. # cache provider credentials
  640. provider_credentials_cache.set(credentials=current_using_credentials)
  641. else:
  642. current_using_credentials = cached_provider_credentials
  643. else:
  644. current_using_credentials = {}
  645. quota_configurations = []
  646. return SystemConfiguration(
  647. enabled=True,
  648. current_quota_type=current_quota_type,
  649. quota_configurations=quota_configurations,
  650. credentials=current_using_credentials,
  651. )
  652. @staticmethod
  653. def _choice_current_using_quota_type(quota_configurations: list[QuotaConfiguration]) -> ProviderQuotaType:
  654. """
  655. Choice current using quota type.
  656. paid quotas > provider free quotas > hosting trial quotas
  657. If there is still quota for the corresponding quota type according to the sorting,
  658. :param quota_configurations:
  659. :return:
  660. """
  661. # convert to dict
  662. quota_type_to_quota_configuration_dict = {
  663. quota_configuration.quota_type: quota_configuration for quota_configuration in quota_configurations
  664. }
  665. last_quota_configuration = None
  666. for quota_type in [ProviderQuotaType.PAID, ProviderQuotaType.FREE, ProviderQuotaType.TRIAL]:
  667. if quota_type in quota_type_to_quota_configuration_dict:
  668. last_quota_configuration = quota_type_to_quota_configuration_dict[quota_type]
  669. if last_quota_configuration.is_valid:
  670. return quota_type
  671. if last_quota_configuration:
  672. return last_quota_configuration.quota_type
  673. raise ValueError("No quota type available")
  674. @staticmethod
  675. def _extract_secret_variables(credential_form_schemas: list[CredentialFormSchema]) -> list[str]:
  676. """
  677. Extract secret input form variables.
  678. :param credential_form_schemas:
  679. :return:
  680. """
  681. secret_input_form_variables = []
  682. for credential_form_schema in credential_form_schemas:
  683. if credential_form_schema.type == FormType.SECRET_INPUT:
  684. secret_input_form_variables.append(credential_form_schema.variable)
  685. return secret_input_form_variables
  686. def _to_model_settings(
  687. self,
  688. provider_entity: ProviderEntity,
  689. provider_model_settings: Optional[list[ProviderModelSetting]] = None,
  690. load_balancing_model_configs: Optional[list[LoadBalancingModelConfig]] = None,
  691. ) -> list[ModelSettings]:
  692. """
  693. Convert to model settings.
  694. :param provider_entity: provider entity
  695. :param provider_model_settings: provider model settings include enabled, load balancing enabled
  696. :param load_balancing_model_configs: load balancing model configs
  697. :return:
  698. """
  699. # Get provider model credential secret variables
  700. model_credential_secret_variables = self._extract_secret_variables(
  701. provider_entity.model_credential_schema.credential_form_schemas
  702. if provider_entity.model_credential_schema
  703. else []
  704. )
  705. model_settings = []
  706. if not provider_model_settings:
  707. return model_settings
  708. for provider_model_setting in provider_model_settings:
  709. load_balancing_configs = []
  710. if provider_model_setting.load_balancing_enabled and load_balancing_model_configs:
  711. for load_balancing_model_config in load_balancing_model_configs:
  712. if (
  713. load_balancing_model_config.model_name == provider_model_setting.model_name
  714. and load_balancing_model_config.model_type == provider_model_setting.model_type
  715. ):
  716. if not load_balancing_model_config.enabled:
  717. continue
  718. if not load_balancing_model_config.encrypted_config:
  719. if load_balancing_model_config.name == "__inherit__":
  720. load_balancing_configs.append(
  721. ModelLoadBalancingConfiguration(
  722. id=load_balancing_model_config.id,
  723. name=load_balancing_model_config.name,
  724. credentials={},
  725. )
  726. )
  727. continue
  728. provider_model_credentials_cache = ProviderCredentialsCache(
  729. tenant_id=load_balancing_model_config.tenant_id,
  730. identity_id=load_balancing_model_config.id,
  731. cache_type=ProviderCredentialsCacheType.LOAD_BALANCING_MODEL,
  732. )
  733. # Get cached provider model credentials
  734. cached_provider_model_credentials = provider_model_credentials_cache.get()
  735. if not cached_provider_model_credentials:
  736. try:
  737. provider_model_credentials = json.loads(load_balancing_model_config.encrypted_config)
  738. except JSONDecodeError:
  739. continue
  740. # Get decoding rsa key and cipher for decrypting credentials
  741. if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None:
  742. self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(
  743. load_balancing_model_config.tenant_id
  744. )
  745. for variable in model_credential_secret_variables:
  746. if variable in provider_model_credentials:
  747. try:
  748. provider_model_credentials[variable] = encrypter.decrypt_token_with_decoding(
  749. provider_model_credentials.get(variable),
  750. self.decoding_rsa_key,
  751. self.decoding_cipher_rsa,
  752. )
  753. except ValueError:
  754. pass
  755. # cache provider model credentials
  756. provider_model_credentials_cache.set(credentials=provider_model_credentials)
  757. else:
  758. provider_model_credentials = cached_provider_model_credentials
  759. load_balancing_configs.append(
  760. ModelLoadBalancingConfiguration(
  761. id=load_balancing_model_config.id,
  762. name=load_balancing_model_config.name,
  763. credentials=provider_model_credentials,
  764. )
  765. )
  766. model_settings.append(
  767. ModelSettings(
  768. model=provider_model_setting.model_name,
  769. model_type=ModelType.value_of(provider_model_setting.model_type),
  770. enabled=provider_model_setting.enabled,
  771. load_balancing_configs=load_balancing_configs if len(load_balancing_configs) > 1 else [],
  772. )
  773. )
  774. return model_settings