You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

provider_manager.py 48KB

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