provider_manager.py 40 KB

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