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

provider_manager.py 50KB

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