provider_manager.py 40 KB

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