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 44KB

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