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

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