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_configuration.py 79KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868
  1. import json
  2. import logging
  3. import re
  4. from collections import defaultdict
  5. from collections.abc import Iterator, Sequence
  6. from json import JSONDecodeError
  7. from pydantic import BaseModel, ConfigDict, Field
  8. from sqlalchemy import func, select
  9. from sqlalchemy.orm import Session
  10. from constants import HIDDEN_VALUE
  11. from core.entities.model_entities import ModelStatus, ModelWithProviderEntity, SimpleModelProviderEntity
  12. from core.entities.provider_entities import (
  13. CustomConfiguration,
  14. ModelSettings,
  15. SystemConfiguration,
  16. SystemConfigurationStatus,
  17. )
  18. from core.helper import encrypter
  19. from core.helper.model_provider_cache import ProviderCredentialsCache, ProviderCredentialsCacheType
  20. from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
  21. from core.model_runtime.entities.provider_entities import (
  22. ConfigurateMethod,
  23. CredentialFormSchema,
  24. FormType,
  25. ProviderEntity,
  26. )
  27. from core.model_runtime.model_providers.__base.ai_model import AIModel
  28. from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
  29. from core.plugin.entities.plugin import ModelProviderID
  30. from extensions.ext_database import db
  31. from libs.datetime_utils import naive_utc_now
  32. from models.provider import (
  33. LoadBalancingModelConfig,
  34. Provider,
  35. ProviderCredential,
  36. ProviderModel,
  37. ProviderModelCredential,
  38. ProviderModelSetting,
  39. ProviderType,
  40. TenantPreferredModelProvider,
  41. )
  42. from services.enterprise.plugin_manager_service import PluginCredentialType
  43. logger = logging.getLogger(__name__)
  44. original_provider_configurate_methods: dict[str, list[ConfigurateMethod]] = {}
  45. class ProviderConfiguration(BaseModel):
  46. """
  47. Provider configuration entity for managing model provider settings.
  48. This class handles:
  49. - Provider credentials CRUD and switch
  50. - Custom Model credentials CRUD and switch
  51. - System vs custom provider switching
  52. - Load balancing configurations
  53. - Model enablement/disablement
  54. TODO: lots of logic in a BaseModel entity should be separated, the exceptions should be classified
  55. """
  56. tenant_id: str
  57. provider: ProviderEntity
  58. preferred_provider_type: ProviderType
  59. using_provider_type: ProviderType
  60. system_configuration: SystemConfiguration
  61. custom_configuration: CustomConfiguration
  62. model_settings: list[ModelSettings]
  63. # pydantic configs
  64. model_config = ConfigDict(protected_namespaces=())
  65. def __init__(self, **data):
  66. super().__init__(**data)
  67. if self.provider.provider not in original_provider_configurate_methods:
  68. original_provider_configurate_methods[self.provider.provider] = []
  69. for configurate_method in self.provider.configurate_methods:
  70. original_provider_configurate_methods[self.provider.provider].append(configurate_method)
  71. if original_provider_configurate_methods[self.provider.provider] == [ConfigurateMethod.CUSTOMIZABLE_MODEL]:
  72. if (
  73. any(
  74. len(quota_configuration.restrict_models) > 0
  75. for quota_configuration in self.system_configuration.quota_configurations
  76. )
  77. and ConfigurateMethod.PREDEFINED_MODEL not in self.provider.configurate_methods
  78. ):
  79. self.provider.configurate_methods.append(ConfigurateMethod.PREDEFINED_MODEL)
  80. def get_current_credentials(self, model_type: ModelType, model: str) -> dict | None:
  81. """
  82. Get current credentials.
  83. :param model_type: model type
  84. :param model: model name
  85. :return:
  86. """
  87. if self.model_settings:
  88. # check if model is disabled by admin
  89. for model_setting in self.model_settings:
  90. if model_setting.model_type == model_type and model_setting.model == model:
  91. if not model_setting.enabled:
  92. raise ValueError(f"Model {model} is disabled.")
  93. if self.using_provider_type == ProviderType.SYSTEM:
  94. restrict_models = []
  95. for quota_configuration in self.system_configuration.quota_configurations:
  96. if self.system_configuration.current_quota_type != quota_configuration.quota_type:
  97. continue
  98. restrict_models = quota_configuration.restrict_models
  99. copy_credentials = (
  100. self.system_configuration.credentials.copy() if self.system_configuration.credentials else {}
  101. )
  102. if restrict_models:
  103. for restrict_model in restrict_models:
  104. if (
  105. restrict_model.model_type == model_type
  106. and restrict_model.model == model
  107. and restrict_model.base_model_name
  108. ):
  109. copy_credentials["base_model_name"] = restrict_model.base_model_name
  110. return copy_credentials
  111. else:
  112. credentials = None
  113. current_credential_id = None
  114. if self.custom_configuration.models:
  115. for model_configuration in self.custom_configuration.models:
  116. if model_configuration.model_type == model_type and model_configuration.model == model:
  117. credentials = model_configuration.credentials
  118. current_credential_id = model_configuration.current_credential_id
  119. break
  120. if not credentials and self.custom_configuration.provider:
  121. credentials = self.custom_configuration.provider.credentials
  122. current_credential_id = self.custom_configuration.provider.current_credential_id
  123. if current_credential_id:
  124. from core.helper.credential_utils import check_credential_policy_compliance
  125. check_credential_policy_compliance(
  126. credential_id=current_credential_id,
  127. provider=self.provider.provider,
  128. credential_type=PluginCredentialType.MODEL,
  129. )
  130. else:
  131. # no current credential id, check all available credentials
  132. if self.custom_configuration.provider:
  133. for credential_configuration in self.custom_configuration.provider.available_credentials:
  134. from core.helper.credential_utils import check_credential_policy_compliance
  135. check_credential_policy_compliance(
  136. credential_id=credential_configuration.credential_id,
  137. provider=self.provider.provider,
  138. credential_type=PluginCredentialType.MODEL,
  139. )
  140. return credentials
  141. def get_system_configuration_status(self) -> SystemConfigurationStatus | None:
  142. """
  143. Get system configuration status.
  144. :return:
  145. """
  146. if self.system_configuration.enabled is False:
  147. return SystemConfigurationStatus.UNSUPPORTED
  148. current_quota_type = self.system_configuration.current_quota_type
  149. current_quota_configuration = next(
  150. (q for q in self.system_configuration.quota_configurations if q.quota_type == current_quota_type), None
  151. )
  152. if current_quota_configuration is None:
  153. return None
  154. if not current_quota_configuration:
  155. return SystemConfigurationStatus.UNSUPPORTED
  156. return (
  157. SystemConfigurationStatus.ACTIVE
  158. if current_quota_configuration.is_valid
  159. else SystemConfigurationStatus.QUOTA_EXCEEDED
  160. )
  161. def is_custom_configuration_available(self) -> bool:
  162. """
  163. Check custom configuration available.
  164. :return:
  165. """
  166. has_provider_credentials = (
  167. self.custom_configuration.provider is not None
  168. and len(self.custom_configuration.provider.available_credentials) > 0
  169. )
  170. has_model_configurations = len(self.custom_configuration.models) > 0
  171. return has_provider_credentials or has_model_configurations
  172. def _get_provider_record(self, session: Session) -> Provider | None:
  173. """
  174. Get custom provider record.
  175. """
  176. # get provider
  177. model_provider_id = ModelProviderID(self.provider.provider)
  178. provider_names = [self.provider.provider]
  179. if model_provider_id.is_langgenius():
  180. provider_names.append(model_provider_id.provider_name)
  181. stmt = select(Provider).where(
  182. Provider.tenant_id == self.tenant_id,
  183. Provider.provider_type == ProviderType.CUSTOM.value,
  184. Provider.provider_name.in_(provider_names),
  185. )
  186. return session.execute(stmt).scalar_one_or_none()
  187. def _get_specific_provider_credential(self, credential_id: str) -> dict | None:
  188. """
  189. Get a specific provider credential by ID.
  190. :param credential_id: Credential ID
  191. :return:
  192. """
  193. # Extract secret variables from provider credential schema
  194. credential_secret_variables = self.extract_secret_variables(
  195. self.provider.provider_credential_schema.credential_form_schemas
  196. if self.provider.provider_credential_schema
  197. else []
  198. )
  199. with Session(db.engine) as session:
  200. # Prefer the actual provider record name if exists (to handle aliased provider names)
  201. provider_record = self._get_provider_record(session)
  202. provider_name = provider_record.provider_name if provider_record else self.provider.provider
  203. stmt = select(ProviderCredential).where(
  204. ProviderCredential.id == credential_id,
  205. ProviderCredential.tenant_id == self.tenant_id,
  206. ProviderCredential.provider_name == provider_name,
  207. )
  208. credential = session.execute(stmt).scalar_one_or_none()
  209. if not credential or not credential.encrypted_config:
  210. raise ValueError(f"Credential with id {credential_id} not found.")
  211. try:
  212. credentials = json.loads(credential.encrypted_config)
  213. except JSONDecodeError:
  214. credentials = {}
  215. # Decrypt secret variables
  216. for key in credential_secret_variables:
  217. if key in credentials and credentials[key] is not None:
  218. try:
  219. credentials[key] = encrypter.decrypt_token(tenant_id=self.tenant_id, token=credentials[key])
  220. except Exception:
  221. pass
  222. return self.obfuscated_credentials(
  223. credentials=credentials,
  224. credential_form_schemas=self.provider.provider_credential_schema.credential_form_schemas
  225. if self.provider.provider_credential_schema
  226. else [],
  227. )
  228. def _check_provider_credential_name_exists(
  229. self, credential_name: str, session: Session, exclude_id: str | None = None
  230. ) -> bool:
  231. """
  232. not allowed same name when create or update a credential
  233. """
  234. stmt = select(ProviderCredential.id).where(
  235. ProviderCredential.tenant_id == self.tenant_id,
  236. ProviderCredential.provider_name == self.provider.provider,
  237. ProviderCredential.credential_name == credential_name,
  238. )
  239. if exclude_id:
  240. stmt = stmt.where(ProviderCredential.id != exclude_id)
  241. return session.execute(stmt).scalar_one_or_none() is not None
  242. def get_provider_credential(self, credential_id: str | None = None) -> dict | None:
  243. """
  244. Get provider credentials.
  245. :param credential_id: if provided, return the specified credential
  246. :return:
  247. """
  248. if credential_id:
  249. return self._get_specific_provider_credential(credential_id)
  250. # Default behavior: return current active provider credentials
  251. credentials = self.custom_configuration.provider.credentials if self.custom_configuration.provider else {}
  252. return self.obfuscated_credentials(
  253. credentials=credentials,
  254. credential_form_schemas=self.provider.provider_credential_schema.credential_form_schemas
  255. if self.provider.provider_credential_schema
  256. else [],
  257. )
  258. def validate_provider_credentials(self, credentials: dict, credential_id: str = "", session: Session | None = None):
  259. """
  260. Validate custom credentials.
  261. :param credentials: provider credentials
  262. :param credential_id: (Optional)If provided, can use existing credential's hidden api key to validate
  263. :param session: optional database session
  264. :return:
  265. """
  266. def _validate(s: Session):
  267. # Get provider credential secret variables
  268. provider_credential_secret_variables = self.extract_secret_variables(
  269. self.provider.provider_credential_schema.credential_form_schemas
  270. if self.provider.provider_credential_schema
  271. else []
  272. )
  273. if credential_id:
  274. try:
  275. stmt = select(ProviderCredential).where(
  276. ProviderCredential.tenant_id == self.tenant_id,
  277. ProviderCredential.provider_name == self.provider.provider,
  278. ProviderCredential.id == credential_id,
  279. )
  280. credential_record = s.execute(stmt).scalar_one_or_none()
  281. # fix origin data
  282. if credential_record and credential_record.encrypted_config:
  283. if not credential_record.encrypted_config.startswith("{"):
  284. original_credentials = {"openai_api_key": credential_record.encrypted_config}
  285. else:
  286. original_credentials = json.loads(credential_record.encrypted_config)
  287. else:
  288. original_credentials = {}
  289. except JSONDecodeError:
  290. original_credentials = {}
  291. # encrypt credentials
  292. for key, value in credentials.items():
  293. if key in provider_credential_secret_variables:
  294. # if send [__HIDDEN__] in secret input, it will be same as original value
  295. if value == HIDDEN_VALUE and key in original_credentials:
  296. credentials[key] = encrypter.decrypt_token(
  297. tenant_id=self.tenant_id, token=original_credentials[key]
  298. )
  299. model_provider_factory = ModelProviderFactory(self.tenant_id)
  300. validated_credentials = model_provider_factory.provider_credentials_validate(
  301. provider=self.provider.provider, credentials=credentials
  302. )
  303. for key, value in validated_credentials.items():
  304. if key in provider_credential_secret_variables:
  305. validated_credentials[key] = encrypter.encrypt_token(self.tenant_id, value)
  306. return validated_credentials
  307. if session:
  308. return _validate(session)
  309. else:
  310. with Session(db.engine) as new_session:
  311. return _validate(new_session)
  312. def _generate_provider_credential_name(self, session) -> str:
  313. """
  314. Generate a unique credential name for provider.
  315. :return: credential name
  316. """
  317. return self._generate_next_api_key_name(
  318. session=session,
  319. query_factory=lambda: select(ProviderCredential).where(
  320. ProviderCredential.tenant_id == self.tenant_id,
  321. ProviderCredential.provider_name == self.provider.provider,
  322. ),
  323. )
  324. def _generate_custom_model_credential_name(self, model: str, model_type: ModelType, session) -> str:
  325. """
  326. Generate a unique credential name for custom model.
  327. :return: credential name
  328. """
  329. return self._generate_next_api_key_name(
  330. session=session,
  331. query_factory=lambda: select(ProviderModelCredential).where(
  332. ProviderModelCredential.tenant_id == self.tenant_id,
  333. ProviderModelCredential.provider_name == self.provider.provider,
  334. ProviderModelCredential.model_name == model,
  335. ProviderModelCredential.model_type == model_type.to_origin_model_type(),
  336. ),
  337. )
  338. def _generate_next_api_key_name(self, session, query_factory) -> str:
  339. """
  340. Generate next available API KEY name by finding the highest numbered suffix.
  341. :param session: database session
  342. :param query_factory: function that returns the SQLAlchemy query
  343. :return: next available API KEY name
  344. """
  345. try:
  346. stmt = query_factory()
  347. credential_records = session.execute(stmt).scalars().all()
  348. if not credential_records:
  349. return "API KEY 1"
  350. # Extract numbers from API KEY pattern using list comprehension
  351. pattern = re.compile(r"^API KEY\s+(\d+)$")
  352. numbers = [
  353. int(match.group(1))
  354. for cr in credential_records
  355. if cr.credential_name and (match := pattern.match(cr.credential_name.strip()))
  356. ]
  357. # Return next sequential number
  358. next_number = max(numbers, default=0) + 1
  359. return f"API KEY {next_number}"
  360. except Exception as e:
  361. logger.warning("Error generating next credential name: %s", str(e))
  362. return "API KEY 1"
  363. def create_provider_credential(self, credentials: dict, credential_name: str | None):
  364. """
  365. Add custom provider credentials.
  366. :param credentials: provider credentials
  367. :param credential_name: credential name
  368. :return:
  369. """
  370. with Session(db.engine) as session:
  371. if credential_name:
  372. if self._check_provider_credential_name_exists(credential_name=credential_name, session=session):
  373. raise ValueError(f"Credential with name '{credential_name}' already exists.")
  374. else:
  375. credential_name = self._generate_provider_credential_name(session)
  376. credentials = self.validate_provider_credentials(credentials=credentials, session=session)
  377. provider_record = self._get_provider_record(session)
  378. try:
  379. new_record = ProviderCredential(
  380. tenant_id=self.tenant_id,
  381. provider_name=self.provider.provider,
  382. encrypted_config=json.dumps(credentials),
  383. credential_name=credential_name,
  384. )
  385. session.add(new_record)
  386. session.flush()
  387. if not provider_record:
  388. # If provider record does not exist, create it
  389. provider_record = Provider(
  390. tenant_id=self.tenant_id,
  391. provider_name=self.provider.provider,
  392. provider_type=ProviderType.CUSTOM.value,
  393. is_valid=True,
  394. credential_id=new_record.id,
  395. )
  396. session.add(provider_record)
  397. provider_model_credentials_cache = ProviderCredentialsCache(
  398. tenant_id=self.tenant_id,
  399. identity_id=provider_record.id,
  400. cache_type=ProviderCredentialsCacheType.PROVIDER,
  401. )
  402. provider_model_credentials_cache.delete()
  403. self.switch_preferred_provider_type(provider_type=ProviderType.CUSTOM, session=session)
  404. session.commit()
  405. except Exception:
  406. session.rollback()
  407. raise
  408. def update_provider_credential(
  409. self,
  410. credentials: dict,
  411. credential_id: str,
  412. credential_name: str | None,
  413. ):
  414. """
  415. update a saved provider credential (by credential_id).
  416. :param credentials: provider credentials
  417. :param credential_id: credential id
  418. :param credential_name: credential name
  419. :return:
  420. """
  421. with Session(db.engine) as session:
  422. if credential_name and self._check_provider_credential_name_exists(
  423. credential_name=credential_name, session=session, exclude_id=credential_id
  424. ):
  425. raise ValueError(f"Credential with name '{credential_name}' already exists.")
  426. credentials = self.validate_provider_credentials(
  427. credentials=credentials, credential_id=credential_id, session=session
  428. )
  429. provider_record = self._get_provider_record(session)
  430. stmt = select(ProviderCredential).where(
  431. ProviderCredential.id == credential_id,
  432. ProviderCredential.tenant_id == self.tenant_id,
  433. ProviderCredential.provider_name == self.provider.provider,
  434. )
  435. # Get the credential record to update
  436. credential_record = session.execute(stmt).scalar_one_or_none()
  437. if not credential_record:
  438. raise ValueError("Credential record not found.")
  439. try:
  440. # Update credential
  441. credential_record.encrypted_config = json.dumps(credentials)
  442. credential_record.updated_at = naive_utc_now()
  443. if credential_name:
  444. credential_record.credential_name = credential_name
  445. session.commit()
  446. if provider_record and provider_record.credential_id == credential_id:
  447. provider_model_credentials_cache = ProviderCredentialsCache(
  448. tenant_id=self.tenant_id,
  449. identity_id=provider_record.id,
  450. cache_type=ProviderCredentialsCacheType.PROVIDER,
  451. )
  452. provider_model_credentials_cache.delete()
  453. self._update_load_balancing_configs_with_credential(
  454. credential_id=credential_id,
  455. credential_record=credential_record,
  456. credential_source="provider",
  457. session=session,
  458. )
  459. except Exception:
  460. session.rollback()
  461. raise
  462. def _update_load_balancing_configs_with_credential(
  463. self,
  464. credential_id: str,
  465. credential_record: ProviderCredential | ProviderModelCredential,
  466. credential_source: str,
  467. session: Session,
  468. ):
  469. """
  470. Update load balancing configurations that reference the given credential_id.
  471. :param credential_id: credential id
  472. :param credential_record: the encrypted_config and credential_name
  473. :param credential_source: the credential comes from the provider_credential(`provider`)
  474. or the provider_model_credential(`custom_model`)
  475. :param session: the database session
  476. :return:
  477. """
  478. # Find all load balancing configs that use this credential_id
  479. stmt = select(LoadBalancingModelConfig).where(
  480. LoadBalancingModelConfig.tenant_id == self.tenant_id,
  481. LoadBalancingModelConfig.provider_name == self.provider.provider,
  482. LoadBalancingModelConfig.credential_id == credential_id,
  483. LoadBalancingModelConfig.credential_source_type == credential_source,
  484. )
  485. load_balancing_configs = session.execute(stmt).scalars().all()
  486. if not load_balancing_configs:
  487. return
  488. # Update each load balancing config with the new credentials
  489. for lb_config in load_balancing_configs:
  490. # Update the encrypted_config with the new credentials
  491. lb_config.encrypted_config = credential_record.encrypted_config
  492. lb_config.name = credential_record.credential_name
  493. lb_config.updated_at = naive_utc_now()
  494. # Clear cache for this load balancing config
  495. lb_credentials_cache = ProviderCredentialsCache(
  496. tenant_id=self.tenant_id,
  497. identity_id=lb_config.id,
  498. cache_type=ProviderCredentialsCacheType.LOAD_BALANCING_MODEL,
  499. )
  500. lb_credentials_cache.delete()
  501. session.commit()
  502. def delete_provider_credential(self, credential_id: str):
  503. """
  504. Delete a saved provider credential (by credential_id).
  505. :param credential_id: credential id
  506. :return:
  507. """
  508. with Session(db.engine) as session:
  509. stmt = select(ProviderCredential).where(
  510. ProviderCredential.id == credential_id,
  511. ProviderCredential.tenant_id == self.tenant_id,
  512. ProviderCredential.provider_name == self.provider.provider,
  513. )
  514. # Get the credential record to update
  515. credential_record = session.execute(stmt).scalar_one_or_none()
  516. if not credential_record:
  517. raise ValueError("Credential record not found.")
  518. # Check if this credential is used in load balancing configs
  519. lb_stmt = select(LoadBalancingModelConfig).where(
  520. LoadBalancingModelConfig.tenant_id == self.tenant_id,
  521. LoadBalancingModelConfig.provider_name == self.provider.provider,
  522. LoadBalancingModelConfig.credential_id == credential_id,
  523. LoadBalancingModelConfig.credential_source_type == "provider",
  524. )
  525. lb_configs_using_credential = session.execute(lb_stmt).scalars().all()
  526. try:
  527. for lb_config in lb_configs_using_credential:
  528. lb_credentials_cache = ProviderCredentialsCache(
  529. tenant_id=self.tenant_id,
  530. identity_id=lb_config.id,
  531. cache_type=ProviderCredentialsCacheType.LOAD_BALANCING_MODEL,
  532. )
  533. lb_credentials_cache.delete()
  534. session.delete(lb_config)
  535. # Check if this is the currently active credential
  536. provider_record = self._get_provider_record(session)
  537. # Check available credentials count BEFORE deleting
  538. # if this is the last credential, we need to delete the provider record
  539. count_stmt = select(func.count(ProviderCredential.id)).where(
  540. ProviderCredential.tenant_id == self.tenant_id,
  541. ProviderCredential.provider_name == self.provider.provider,
  542. )
  543. available_credentials_count = session.execute(count_stmt).scalar() or 0
  544. session.delete(credential_record)
  545. if provider_record and available_credentials_count <= 1:
  546. # If all credentials are deleted, delete the provider record
  547. session.delete(provider_record)
  548. provider_model_credentials_cache = ProviderCredentialsCache(
  549. tenant_id=self.tenant_id,
  550. identity_id=provider_record.id,
  551. cache_type=ProviderCredentialsCacheType.PROVIDER,
  552. )
  553. provider_model_credentials_cache.delete()
  554. self.switch_preferred_provider_type(provider_type=ProviderType.SYSTEM, session=session)
  555. elif provider_record and provider_record.credential_id == credential_id:
  556. provider_record.credential_id = None
  557. provider_record.updated_at = naive_utc_now()
  558. provider_model_credentials_cache = ProviderCredentialsCache(
  559. tenant_id=self.tenant_id,
  560. identity_id=provider_record.id,
  561. cache_type=ProviderCredentialsCacheType.PROVIDER,
  562. )
  563. provider_model_credentials_cache.delete()
  564. self.switch_preferred_provider_type(provider_type=ProviderType.SYSTEM, session=session)
  565. session.commit()
  566. except Exception:
  567. session.rollback()
  568. raise
  569. def switch_active_provider_credential(self, credential_id: str):
  570. """
  571. Switch active provider credential (copy the selected one into current active snapshot).
  572. :param credential_id: credential id
  573. :return:
  574. """
  575. with Session(db.engine) as session:
  576. stmt = select(ProviderCredential).where(
  577. ProviderCredential.id == credential_id,
  578. ProviderCredential.tenant_id == self.tenant_id,
  579. ProviderCredential.provider_name == self.provider.provider,
  580. )
  581. credential_record = session.execute(stmt).scalar_one_or_none()
  582. if not credential_record:
  583. raise ValueError("Credential record not found.")
  584. provider_record = self._get_provider_record(session)
  585. if not provider_record:
  586. raise ValueError("Provider record not found.")
  587. try:
  588. provider_record.credential_id = credential_record.id
  589. provider_record.updated_at = naive_utc_now()
  590. session.commit()
  591. provider_model_credentials_cache = ProviderCredentialsCache(
  592. tenant_id=self.tenant_id,
  593. identity_id=provider_record.id,
  594. cache_type=ProviderCredentialsCacheType.PROVIDER,
  595. )
  596. provider_model_credentials_cache.delete()
  597. self.switch_preferred_provider_type(ProviderType.CUSTOM, session=session)
  598. except Exception:
  599. session.rollback()
  600. raise
  601. def _get_custom_model_record(
  602. self,
  603. model_type: ModelType,
  604. model: str,
  605. session: Session,
  606. ) -> ProviderModel | None:
  607. """
  608. Get custom model credentials.
  609. """
  610. # get provider model
  611. model_provider_id = ModelProviderID(self.provider.provider)
  612. provider_names = [self.provider.provider]
  613. if model_provider_id.is_langgenius():
  614. provider_names.append(model_provider_id.provider_name)
  615. stmt = select(ProviderModel).where(
  616. ProviderModel.tenant_id == self.tenant_id,
  617. ProviderModel.provider_name.in_(provider_names),
  618. ProviderModel.model_name == model,
  619. ProviderModel.model_type == model_type.to_origin_model_type(),
  620. )
  621. return session.execute(stmt).scalar_one_or_none()
  622. def _get_specific_custom_model_credential(
  623. self, model_type: ModelType, model: str, credential_id: str
  624. ) -> dict | None:
  625. """
  626. Get a specific provider credential by ID.
  627. :param credential_id: Credential ID
  628. :return:
  629. """
  630. model_credential_secret_variables = self.extract_secret_variables(
  631. self.provider.model_credential_schema.credential_form_schemas
  632. if self.provider.model_credential_schema
  633. else []
  634. )
  635. with Session(db.engine) as session:
  636. stmt = select(ProviderModelCredential).where(
  637. ProviderModelCredential.id == credential_id,
  638. ProviderModelCredential.tenant_id == self.tenant_id,
  639. ProviderModelCredential.provider_name == self.provider.provider,
  640. ProviderModelCredential.model_name == model,
  641. ProviderModelCredential.model_type == model_type.to_origin_model_type(),
  642. )
  643. credential_record = session.execute(stmt).scalar_one_or_none()
  644. if not credential_record or not credential_record.encrypted_config:
  645. raise ValueError(f"Credential with id {credential_id} not found.")
  646. try:
  647. credentials = json.loads(credential_record.encrypted_config)
  648. except JSONDecodeError:
  649. credentials = {}
  650. # Decrypt secret variables
  651. for key in model_credential_secret_variables:
  652. if key in credentials and credentials[key] is not None:
  653. try:
  654. credentials[key] = encrypter.decrypt_token(tenant_id=self.tenant_id, token=credentials[key])
  655. except Exception:
  656. pass
  657. current_credential_id = credential_record.id
  658. current_credential_name = credential_record.credential_name
  659. credentials = self.obfuscated_credentials(
  660. credentials=credentials,
  661. credential_form_schemas=self.provider.model_credential_schema.credential_form_schemas
  662. if self.provider.model_credential_schema
  663. else [],
  664. )
  665. return {
  666. "current_credential_id": current_credential_id,
  667. "current_credential_name": current_credential_name,
  668. "credentials": credentials,
  669. }
  670. def _check_custom_model_credential_name_exists(
  671. self, model_type: ModelType, model: str, credential_name: str, session: Session, exclude_id: str | None = None
  672. ) -> bool:
  673. """
  674. not allowed same name when create or update a credential
  675. """
  676. stmt = select(ProviderModelCredential).where(
  677. ProviderModelCredential.tenant_id == self.tenant_id,
  678. ProviderModelCredential.provider_name == self.provider.provider,
  679. ProviderModelCredential.model_name == model,
  680. ProviderModelCredential.model_type == model_type.to_origin_model_type(),
  681. ProviderModelCredential.credential_name == credential_name,
  682. )
  683. if exclude_id:
  684. stmt = stmt.where(ProviderModelCredential.id != exclude_id)
  685. return session.execute(stmt).scalar_one_or_none() is not None
  686. def get_custom_model_credential(self, model_type: ModelType, model: str, credential_id: str | None) -> dict | None:
  687. """
  688. Get custom model credentials.
  689. :param model_type: model type
  690. :param model: model name
  691. :return:
  692. """
  693. # If credential_id is provided, return the specific credential
  694. if credential_id:
  695. return self._get_specific_custom_model_credential(
  696. model_type=model_type, model=model, credential_id=credential_id
  697. )
  698. for model_configuration in self.custom_configuration.models:
  699. if (
  700. model_configuration.model_type == model_type
  701. and model_configuration.model == model
  702. and model_configuration.credentials
  703. ):
  704. current_credential_id = model_configuration.current_credential_id
  705. current_credential_name = model_configuration.current_credential_name
  706. credentials = self.obfuscated_credentials(
  707. credentials=model_configuration.credentials,
  708. credential_form_schemas=self.provider.model_credential_schema.credential_form_schemas
  709. if self.provider.model_credential_schema
  710. else [],
  711. )
  712. return {
  713. "current_credential_id": current_credential_id,
  714. "current_credential_name": current_credential_name,
  715. "credentials": credentials,
  716. }
  717. return None
  718. def validate_custom_model_credentials(
  719. self,
  720. model_type: ModelType,
  721. model: str,
  722. credentials: dict,
  723. credential_id: str = "",
  724. session: Session | None = None,
  725. ):
  726. """
  727. Validate custom model credentials.
  728. :param model_type: model type
  729. :param model: model name
  730. :param credentials: model credentials dict
  731. :param credential_id: (Optional)If provided, can use existing credential's hidden api key to validate
  732. :return:
  733. """
  734. def _validate(s: Session):
  735. # Get provider credential secret variables
  736. provider_credential_secret_variables = self.extract_secret_variables(
  737. self.provider.model_credential_schema.credential_form_schemas
  738. if self.provider.model_credential_schema
  739. else []
  740. )
  741. if credential_id:
  742. try:
  743. stmt = select(ProviderModelCredential).where(
  744. ProviderModelCredential.id == credential_id,
  745. ProviderModelCredential.tenant_id == self.tenant_id,
  746. ProviderModelCredential.provider_name == self.provider.provider,
  747. ProviderModelCredential.model_name == model,
  748. ProviderModelCredential.model_type == model_type.to_origin_model_type(),
  749. )
  750. credential_record = s.execute(stmt).scalar_one_or_none()
  751. original_credentials = (
  752. json.loads(credential_record.encrypted_config)
  753. if credential_record and credential_record.encrypted_config
  754. else {}
  755. )
  756. except JSONDecodeError:
  757. original_credentials = {}
  758. # decrypt credentials
  759. for key, value in credentials.items():
  760. if key in provider_credential_secret_variables:
  761. # if send [__HIDDEN__] in secret input, it will be same as original value
  762. if value == HIDDEN_VALUE and key in original_credentials:
  763. credentials[key] = encrypter.decrypt_token(
  764. tenant_id=self.tenant_id, token=original_credentials[key]
  765. )
  766. model_provider_factory = ModelProviderFactory(self.tenant_id)
  767. validated_credentials = model_provider_factory.model_credentials_validate(
  768. provider=self.provider.provider, model_type=model_type, model=model, credentials=credentials
  769. )
  770. for key, value in validated_credentials.items():
  771. if key in provider_credential_secret_variables:
  772. validated_credentials[key] = encrypter.encrypt_token(self.tenant_id, value)
  773. return validated_credentials
  774. if session:
  775. return _validate(session)
  776. else:
  777. with Session(db.engine) as new_session:
  778. return _validate(new_session)
  779. def create_custom_model_credential(
  780. self, model_type: ModelType, model: str, credentials: dict, credential_name: str | None
  781. ) -> None:
  782. """
  783. Create a custom model credential.
  784. :param model_type: model type
  785. :param model: model name
  786. :param credentials: model credentials dict
  787. :return:
  788. """
  789. with Session(db.engine) as session:
  790. if credential_name:
  791. if self._check_custom_model_credential_name_exists(
  792. model=model, model_type=model_type, credential_name=credential_name, session=session
  793. ):
  794. raise ValueError(f"Model credential with name '{credential_name}' already exists for {model}.")
  795. else:
  796. credential_name = self._generate_custom_model_credential_name(
  797. model=model, model_type=model_type, session=session
  798. )
  799. # validate custom model config
  800. credentials = self.validate_custom_model_credentials(
  801. model_type=model_type, model=model, credentials=credentials, session=session
  802. )
  803. provider_model_record = self._get_custom_model_record(model_type=model_type, model=model, session=session)
  804. try:
  805. credential = ProviderModelCredential(
  806. tenant_id=self.tenant_id,
  807. provider_name=self.provider.provider,
  808. model_name=model,
  809. model_type=model_type.to_origin_model_type(),
  810. encrypted_config=json.dumps(credentials),
  811. credential_name=credential_name,
  812. )
  813. session.add(credential)
  814. session.flush()
  815. # save provider model
  816. if not provider_model_record:
  817. provider_model_record = ProviderModel(
  818. tenant_id=self.tenant_id,
  819. provider_name=self.provider.provider,
  820. model_name=model,
  821. model_type=model_type.to_origin_model_type(),
  822. credential_id=credential.id,
  823. is_valid=True,
  824. )
  825. session.add(provider_model_record)
  826. session.commit()
  827. provider_model_credentials_cache = ProviderCredentialsCache(
  828. tenant_id=self.tenant_id,
  829. identity_id=provider_model_record.id,
  830. cache_type=ProviderCredentialsCacheType.MODEL,
  831. )
  832. provider_model_credentials_cache.delete()
  833. except Exception:
  834. session.rollback()
  835. raise
  836. def update_custom_model_credential(
  837. self, model_type: ModelType, model: str, credentials: dict, credential_name: str | None, credential_id: str
  838. ) -> None:
  839. """
  840. Update a custom model credential.
  841. :param model_type: model type
  842. :param model: model name
  843. :param credentials: model credentials dict
  844. :param credential_name: credential name
  845. :param credential_id: credential id
  846. :return:
  847. """
  848. with Session(db.engine) as session:
  849. if credential_name and self._check_custom_model_credential_name_exists(
  850. model=model,
  851. model_type=model_type,
  852. credential_name=credential_name,
  853. session=session,
  854. exclude_id=credential_id,
  855. ):
  856. raise ValueError(f"Model credential with name '{credential_name}' already exists for {model}.")
  857. # validate custom model config
  858. credentials = self.validate_custom_model_credentials(
  859. model_type=model_type,
  860. model=model,
  861. credentials=credentials,
  862. credential_id=credential_id,
  863. session=session,
  864. )
  865. provider_model_record = self._get_custom_model_record(model_type=model_type, model=model, session=session)
  866. stmt = select(ProviderModelCredential).where(
  867. ProviderModelCredential.id == credential_id,
  868. ProviderModelCredential.tenant_id == self.tenant_id,
  869. ProviderModelCredential.provider_name == self.provider.provider,
  870. ProviderModelCredential.model_name == model,
  871. ProviderModelCredential.model_type == model_type.to_origin_model_type(),
  872. )
  873. credential_record = session.execute(stmt).scalar_one_or_none()
  874. if not credential_record:
  875. raise ValueError("Credential record not found.")
  876. try:
  877. # Update credential
  878. credential_record.encrypted_config = json.dumps(credentials)
  879. credential_record.updated_at = naive_utc_now()
  880. if credential_name:
  881. credential_record.credential_name = credential_name
  882. session.commit()
  883. if provider_model_record and provider_model_record.credential_id == credential_id:
  884. provider_model_credentials_cache = ProviderCredentialsCache(
  885. tenant_id=self.tenant_id,
  886. identity_id=provider_model_record.id,
  887. cache_type=ProviderCredentialsCacheType.MODEL,
  888. )
  889. provider_model_credentials_cache.delete()
  890. self._update_load_balancing_configs_with_credential(
  891. credential_id=credential_id,
  892. credential_record=credential_record,
  893. credential_source="custom_model",
  894. session=session,
  895. )
  896. except Exception:
  897. session.rollback()
  898. raise
  899. def delete_custom_model_credential(self, model_type: ModelType, model: str, credential_id: str):
  900. """
  901. Delete a saved provider credential (by credential_id).
  902. :param credential_id: credential id
  903. :return:
  904. """
  905. with Session(db.engine) as session:
  906. stmt = select(ProviderModelCredential).where(
  907. ProviderModelCredential.id == credential_id,
  908. ProviderModelCredential.tenant_id == self.tenant_id,
  909. ProviderModelCredential.provider_name == self.provider.provider,
  910. ProviderModelCredential.model_name == model,
  911. ProviderModelCredential.model_type == model_type.to_origin_model_type(),
  912. )
  913. credential_record = session.execute(stmt).scalar_one_or_none()
  914. if not credential_record:
  915. raise ValueError("Credential record not found.")
  916. lb_stmt = select(LoadBalancingModelConfig).where(
  917. LoadBalancingModelConfig.tenant_id == self.tenant_id,
  918. LoadBalancingModelConfig.provider_name == self.provider.provider,
  919. LoadBalancingModelConfig.credential_id == credential_id,
  920. LoadBalancingModelConfig.credential_source_type == "custom_model",
  921. )
  922. lb_configs_using_credential = session.execute(lb_stmt).scalars().all()
  923. try:
  924. for lb_config in lb_configs_using_credential:
  925. lb_credentials_cache = ProviderCredentialsCache(
  926. tenant_id=self.tenant_id,
  927. identity_id=lb_config.id,
  928. cache_type=ProviderCredentialsCacheType.LOAD_BALANCING_MODEL,
  929. )
  930. lb_credentials_cache.delete()
  931. session.delete(lb_config)
  932. # Check if this is the currently active credential
  933. provider_model_record = self._get_custom_model_record(model_type, model, session=session)
  934. # Check available credentials count BEFORE deleting
  935. # if this is the last credential, we need to delete the custom model record
  936. count_stmt = select(func.count(ProviderModelCredential.id)).where(
  937. ProviderModelCredential.tenant_id == self.tenant_id,
  938. ProviderModelCredential.provider_name == self.provider.provider,
  939. ProviderModelCredential.model_name == model,
  940. ProviderModelCredential.model_type == model_type.to_origin_model_type(),
  941. )
  942. available_credentials_count = session.execute(count_stmt).scalar() or 0
  943. session.delete(credential_record)
  944. if provider_model_record and available_credentials_count <= 1:
  945. # If all credentials are deleted, delete the custom model record
  946. session.delete(provider_model_record)
  947. elif provider_model_record and provider_model_record.credential_id == credential_id:
  948. provider_model_record.credential_id = None
  949. provider_model_record.updated_at = naive_utc_now()
  950. provider_model_credentials_cache = ProviderCredentialsCache(
  951. tenant_id=self.tenant_id,
  952. identity_id=provider_model_record.id,
  953. cache_type=ProviderCredentialsCacheType.PROVIDER,
  954. )
  955. provider_model_credentials_cache.delete()
  956. session.commit()
  957. except Exception:
  958. session.rollback()
  959. raise
  960. def add_model_credential_to_model(self, model_type: ModelType, model: str, credential_id: str):
  961. """
  962. if model list exist this custom model, switch the custom model credential.
  963. if model list not exist this custom model, use the credential to add a new custom model record.
  964. :param model_type: model type
  965. :param model: model name
  966. :param credential_id: credential id
  967. :return:
  968. """
  969. with Session(db.engine) as session:
  970. stmt = select(ProviderModelCredential).where(
  971. ProviderModelCredential.id == credential_id,
  972. ProviderModelCredential.tenant_id == self.tenant_id,
  973. ProviderModelCredential.provider_name == self.provider.provider,
  974. ProviderModelCredential.model_name == model,
  975. ProviderModelCredential.model_type == model_type.to_origin_model_type(),
  976. )
  977. credential_record = session.execute(stmt).scalar_one_or_none()
  978. if not credential_record:
  979. raise ValueError("Credential record not found.")
  980. # validate custom model config
  981. provider_model_record = self._get_custom_model_record(model_type=model_type, model=model, session=session)
  982. if not provider_model_record:
  983. # create provider model record
  984. provider_model_record = ProviderModel(
  985. tenant_id=self.tenant_id,
  986. provider_name=self.provider.provider,
  987. model_name=model,
  988. model_type=model_type.to_origin_model_type(),
  989. is_valid=True,
  990. credential_id=credential_id,
  991. )
  992. else:
  993. if provider_model_record.credential_id == credential_record.id:
  994. raise ValueError("Can't add same credential")
  995. provider_model_record.credential_id = credential_record.id
  996. provider_model_record.updated_at = naive_utc_now()
  997. session.add(provider_model_record)
  998. session.commit()
  999. def switch_custom_model_credential(self, model_type: ModelType, model: str, credential_id: str):
  1000. """
  1001. switch the custom model credential.
  1002. :param model_type: model type
  1003. :param model: model name
  1004. :param credential_id: credential id
  1005. :return:
  1006. """
  1007. with Session(db.engine) as session:
  1008. stmt = select(ProviderModelCredential).where(
  1009. ProviderModelCredential.id == credential_id,
  1010. ProviderModelCredential.tenant_id == self.tenant_id,
  1011. ProviderModelCredential.provider_name == self.provider.provider,
  1012. ProviderModelCredential.model_name == model,
  1013. ProviderModelCredential.model_type == model_type.to_origin_model_type(),
  1014. )
  1015. credential_record = session.execute(stmt).scalar_one_or_none()
  1016. if not credential_record:
  1017. raise ValueError("Credential record not found.")
  1018. provider_model_record = self._get_custom_model_record(model_type=model_type, model=model, session=session)
  1019. if not provider_model_record:
  1020. raise ValueError("The custom model record not found.")
  1021. provider_model_record.credential_id = credential_record.id
  1022. provider_model_record.updated_at = naive_utc_now()
  1023. session.add(provider_model_record)
  1024. session.commit()
  1025. def delete_custom_model(self, model_type: ModelType, model: str):
  1026. """
  1027. Delete custom model.
  1028. :param model_type: model type
  1029. :param model: model name
  1030. :return:
  1031. """
  1032. with Session(db.engine) as session:
  1033. # get provider model
  1034. provider_model_record = self._get_custom_model_record(model_type=model_type, model=model, session=session)
  1035. # delete provider model
  1036. if provider_model_record:
  1037. session.delete(provider_model_record)
  1038. session.commit()
  1039. provider_model_credentials_cache = ProviderCredentialsCache(
  1040. tenant_id=self.tenant_id,
  1041. identity_id=provider_model_record.id,
  1042. cache_type=ProviderCredentialsCacheType.MODEL,
  1043. )
  1044. provider_model_credentials_cache.delete()
  1045. def _get_provider_model_setting(
  1046. self, model_type: ModelType, model: str, session: Session
  1047. ) -> ProviderModelSetting | None:
  1048. """
  1049. Get provider model setting.
  1050. """
  1051. model_provider_id = ModelProviderID(self.provider.provider)
  1052. provider_names = [self.provider.provider]
  1053. if model_provider_id.is_langgenius():
  1054. provider_names.append(model_provider_id.provider_name)
  1055. stmt = select(ProviderModelSetting).where(
  1056. ProviderModelSetting.tenant_id == self.tenant_id,
  1057. ProviderModelSetting.provider_name.in_(provider_names),
  1058. ProviderModelSetting.model_type == model_type.to_origin_model_type(),
  1059. ProviderModelSetting.model_name == model,
  1060. )
  1061. return session.execute(stmt).scalars().first()
  1062. def enable_model(self, model_type: ModelType, model: str) -> ProviderModelSetting:
  1063. """
  1064. Enable model.
  1065. :param model_type: model type
  1066. :param model: model name
  1067. :return:
  1068. """
  1069. with Session(db.engine) as session:
  1070. model_setting = self._get_provider_model_setting(model_type=model_type, model=model, session=session)
  1071. if model_setting:
  1072. model_setting.enabled = True
  1073. model_setting.updated_at = naive_utc_now()
  1074. else:
  1075. model_setting = ProviderModelSetting(
  1076. tenant_id=self.tenant_id,
  1077. provider_name=self.provider.provider,
  1078. model_type=model_type.to_origin_model_type(),
  1079. model_name=model,
  1080. enabled=True,
  1081. )
  1082. session.add(model_setting)
  1083. session.commit()
  1084. return model_setting
  1085. def disable_model(self, model_type: ModelType, model: str) -> ProviderModelSetting:
  1086. """
  1087. Disable model.
  1088. :param model_type: model type
  1089. :param model: model name
  1090. :return:
  1091. """
  1092. with Session(db.engine) as session:
  1093. model_setting = self._get_provider_model_setting(model_type=model_type, model=model, session=session)
  1094. if model_setting:
  1095. model_setting.enabled = False
  1096. model_setting.updated_at = naive_utc_now()
  1097. else:
  1098. model_setting = ProviderModelSetting(
  1099. tenant_id=self.tenant_id,
  1100. provider_name=self.provider.provider,
  1101. model_type=model_type.to_origin_model_type(),
  1102. model_name=model,
  1103. enabled=False,
  1104. )
  1105. session.add(model_setting)
  1106. session.commit()
  1107. return model_setting
  1108. def get_provider_model_setting(self, model_type: ModelType, model: str) -> ProviderModelSetting | None:
  1109. """
  1110. Get provider model setting.
  1111. :param model_type: model type
  1112. :param model: model name
  1113. :return:
  1114. """
  1115. with Session(db.engine) as session:
  1116. return self._get_provider_model_setting(model_type=model_type, model=model, session=session)
  1117. def enable_model_load_balancing(self, model_type: ModelType, model: str) -> ProviderModelSetting:
  1118. """
  1119. Enable model load balancing.
  1120. :param model_type: model type
  1121. :param model: model name
  1122. :return:
  1123. """
  1124. model_provider_id = ModelProviderID(self.provider.provider)
  1125. provider_names = [self.provider.provider]
  1126. if model_provider_id.is_langgenius():
  1127. provider_names.append(model_provider_id.provider_name)
  1128. with Session(db.engine) as session:
  1129. stmt = select(func.count(LoadBalancingModelConfig.id)).where(
  1130. LoadBalancingModelConfig.tenant_id == self.tenant_id,
  1131. LoadBalancingModelConfig.provider_name.in_(provider_names),
  1132. LoadBalancingModelConfig.model_type == model_type.to_origin_model_type(),
  1133. LoadBalancingModelConfig.model_name == model,
  1134. )
  1135. load_balancing_config_count = session.execute(stmt).scalar() or 0
  1136. if load_balancing_config_count <= 1:
  1137. raise ValueError("Model load balancing configuration must be more than 1.")
  1138. model_setting = self._get_provider_model_setting(model_type=model_type, model=model, session=session)
  1139. if model_setting:
  1140. model_setting.load_balancing_enabled = True
  1141. model_setting.updated_at = naive_utc_now()
  1142. else:
  1143. model_setting = ProviderModelSetting(
  1144. tenant_id=self.tenant_id,
  1145. provider_name=self.provider.provider,
  1146. model_type=model_type.to_origin_model_type(),
  1147. model_name=model,
  1148. load_balancing_enabled=True,
  1149. )
  1150. session.add(model_setting)
  1151. session.commit()
  1152. return model_setting
  1153. def disable_model_load_balancing(self, model_type: ModelType, model: str) -> ProviderModelSetting:
  1154. """
  1155. Disable model load balancing.
  1156. :param model_type: model type
  1157. :param model: model name
  1158. :return:
  1159. """
  1160. with Session(db.engine) as session:
  1161. model_setting = self._get_provider_model_setting(model_type=model_type, model=model, session=session)
  1162. if model_setting:
  1163. model_setting.load_balancing_enabled = False
  1164. model_setting.updated_at = naive_utc_now()
  1165. else:
  1166. model_setting = ProviderModelSetting(
  1167. tenant_id=self.tenant_id,
  1168. provider_name=self.provider.provider,
  1169. model_type=model_type.to_origin_model_type(),
  1170. model_name=model,
  1171. load_balancing_enabled=False,
  1172. )
  1173. session.add(model_setting)
  1174. session.commit()
  1175. return model_setting
  1176. def get_model_type_instance(self, model_type: ModelType) -> AIModel:
  1177. """
  1178. Get current model type instance.
  1179. :param model_type: model type
  1180. :return:
  1181. """
  1182. model_provider_factory = ModelProviderFactory(self.tenant_id)
  1183. # Get model instance of LLM
  1184. return model_provider_factory.get_model_type_instance(provider=self.provider.provider, model_type=model_type)
  1185. def get_model_schema(self, model_type: ModelType, model: str, credentials: dict | None) -> AIModelEntity | None:
  1186. """
  1187. Get model schema
  1188. """
  1189. model_provider_factory = ModelProviderFactory(self.tenant_id)
  1190. return model_provider_factory.get_model_schema(
  1191. provider=self.provider.provider, model_type=model_type, model=model, credentials=credentials
  1192. )
  1193. def switch_preferred_provider_type(self, provider_type: ProviderType, session: Session | None = None):
  1194. """
  1195. Switch preferred provider type.
  1196. :param provider_type:
  1197. :return:
  1198. """
  1199. if provider_type == self.preferred_provider_type:
  1200. return
  1201. if provider_type == ProviderType.SYSTEM and not self.system_configuration.enabled:
  1202. return
  1203. def _switch(s: Session):
  1204. # get preferred provider
  1205. model_provider_id = ModelProviderID(self.provider.provider)
  1206. provider_names = [self.provider.provider]
  1207. if model_provider_id.is_langgenius():
  1208. provider_names.append(model_provider_id.provider_name)
  1209. stmt = select(TenantPreferredModelProvider).where(
  1210. TenantPreferredModelProvider.tenant_id == self.tenant_id,
  1211. TenantPreferredModelProvider.provider_name.in_(provider_names),
  1212. )
  1213. preferred_model_provider = s.execute(stmt).scalars().first()
  1214. if preferred_model_provider:
  1215. preferred_model_provider.preferred_provider_type = provider_type.value
  1216. else:
  1217. preferred_model_provider = TenantPreferredModelProvider(
  1218. tenant_id=self.tenant_id,
  1219. provider_name=self.provider.provider,
  1220. preferred_provider_type=provider_type.value,
  1221. )
  1222. s.add(preferred_model_provider)
  1223. s.commit()
  1224. if session:
  1225. return _switch(session)
  1226. else:
  1227. with Session(db.engine) as session:
  1228. return _switch(session)
  1229. def extract_secret_variables(self, credential_form_schemas: list[CredentialFormSchema]) -> list[str]:
  1230. """
  1231. Extract secret input form variables.
  1232. :param credential_form_schemas:
  1233. :return:
  1234. """
  1235. secret_input_form_variables = []
  1236. for credential_form_schema in credential_form_schemas:
  1237. if credential_form_schema.type == FormType.SECRET_INPUT:
  1238. secret_input_form_variables.append(credential_form_schema.variable)
  1239. return secret_input_form_variables
  1240. def obfuscated_credentials(self, credentials: dict, credential_form_schemas: list[CredentialFormSchema]):
  1241. """
  1242. Obfuscated credentials.
  1243. :param credentials: credentials
  1244. :param credential_form_schemas: credential form schemas
  1245. :return:
  1246. """
  1247. # Get provider credential secret variables
  1248. credential_secret_variables = self.extract_secret_variables(credential_form_schemas)
  1249. # Obfuscate provider credentials
  1250. copy_credentials = credentials.copy()
  1251. for key, value in copy_credentials.items():
  1252. if key in credential_secret_variables:
  1253. copy_credentials[key] = encrypter.obfuscated_token(value)
  1254. return copy_credentials
  1255. def get_provider_model(
  1256. self, model_type: ModelType, model: str, only_active: bool = False
  1257. ) -> ModelWithProviderEntity | None:
  1258. """
  1259. Get provider model.
  1260. :param model_type: model type
  1261. :param model: model name
  1262. :param only_active: return active model only
  1263. :return:
  1264. """
  1265. provider_models = self.get_provider_models(model_type, only_active, model)
  1266. for provider_model in provider_models:
  1267. if provider_model.model == model:
  1268. return provider_model
  1269. return None
  1270. def get_provider_models(
  1271. self, model_type: ModelType | None = None, only_active: bool = False, model: str | None = None
  1272. ) -> list[ModelWithProviderEntity]:
  1273. """
  1274. Get provider models.
  1275. :param model_type: model type
  1276. :param only_active: only active models
  1277. :param model: model name
  1278. :return:
  1279. """
  1280. model_provider_factory = ModelProviderFactory(self.tenant_id)
  1281. provider_schema = model_provider_factory.get_provider_schema(self.provider.provider)
  1282. model_types: list[ModelType] = []
  1283. if model_type:
  1284. model_types.append(model_type)
  1285. else:
  1286. model_types = list(provider_schema.supported_model_types)
  1287. # Group model settings by model type and model
  1288. model_setting_map: defaultdict[ModelType, dict[str, ModelSettings]] = defaultdict(dict)
  1289. for model_setting in self.model_settings:
  1290. model_setting_map[model_setting.model_type][model_setting.model] = model_setting
  1291. if self.using_provider_type == ProviderType.SYSTEM:
  1292. provider_models = self._get_system_provider_models(
  1293. model_types=model_types, provider_schema=provider_schema, model_setting_map=model_setting_map
  1294. )
  1295. else:
  1296. provider_models = self._get_custom_provider_models(
  1297. model_types=model_types,
  1298. provider_schema=provider_schema,
  1299. model_setting_map=model_setting_map,
  1300. model=model,
  1301. )
  1302. if only_active:
  1303. provider_models = [m for m in provider_models if m.status == ModelStatus.ACTIVE]
  1304. # resort provider_models
  1305. # Optimize sorting logic: first sort by provider.position order, then by model_type.value
  1306. # Get the position list for model types (retrieve only once for better performance)
  1307. model_type_positions = {}
  1308. if hasattr(self.provider, "position") and self.provider.position:
  1309. model_type_positions = self.provider.position
  1310. def get_sort_key(model: ModelWithProviderEntity):
  1311. # Get the position list for the current model type
  1312. positions = model_type_positions.get(model.model_type.value, [])
  1313. # If the model name is in the position list, use its index for sorting
  1314. # Otherwise use a large value (list length) to place undefined models at the end
  1315. position_index = positions.index(model.model) if model.model in positions else len(positions)
  1316. # Return composite sort key: (model_type value, model position index)
  1317. return (model.model_type.value, position_index)
  1318. # Sort using the composite sort key
  1319. return sorted(provider_models, key=get_sort_key)
  1320. def _get_system_provider_models(
  1321. self,
  1322. model_types: Sequence[ModelType],
  1323. provider_schema: ProviderEntity,
  1324. model_setting_map: dict[ModelType, dict[str, ModelSettings]],
  1325. ) -> list[ModelWithProviderEntity]:
  1326. """
  1327. Get system provider models.
  1328. :param model_types: model types
  1329. :param provider_schema: provider schema
  1330. :param model_setting_map: model setting map
  1331. :return:
  1332. """
  1333. provider_models = []
  1334. for model_type in model_types:
  1335. for m in provider_schema.models:
  1336. if m.model_type != model_type:
  1337. continue
  1338. status = ModelStatus.ACTIVE
  1339. if m.model_type in model_setting_map and m.model in model_setting_map[m.model_type]:
  1340. model_setting = model_setting_map[m.model_type][m.model]
  1341. if model_setting.enabled is False:
  1342. status = ModelStatus.DISABLED
  1343. provider_models.append(
  1344. ModelWithProviderEntity(
  1345. model=m.model,
  1346. label=m.label,
  1347. model_type=m.model_type,
  1348. features=m.features,
  1349. fetch_from=m.fetch_from,
  1350. model_properties=m.model_properties,
  1351. deprecated=m.deprecated,
  1352. provider=SimpleModelProviderEntity(self.provider),
  1353. status=status,
  1354. )
  1355. )
  1356. if self.provider.provider not in original_provider_configurate_methods:
  1357. original_provider_configurate_methods[self.provider.provider] = []
  1358. for configurate_method in provider_schema.configurate_methods:
  1359. original_provider_configurate_methods[self.provider.provider].append(configurate_method)
  1360. should_use_custom_model = False
  1361. if original_provider_configurate_methods[self.provider.provider] == [ConfigurateMethod.CUSTOMIZABLE_MODEL]:
  1362. should_use_custom_model = True
  1363. for quota_configuration in self.system_configuration.quota_configurations:
  1364. if self.system_configuration.current_quota_type != quota_configuration.quota_type:
  1365. continue
  1366. restrict_models = quota_configuration.restrict_models
  1367. if len(restrict_models) == 0:
  1368. break
  1369. if should_use_custom_model:
  1370. if original_provider_configurate_methods[self.provider.provider] == [
  1371. ConfigurateMethod.CUSTOMIZABLE_MODEL
  1372. ]:
  1373. # only customizable model
  1374. for restrict_model in restrict_models:
  1375. copy_credentials = (
  1376. self.system_configuration.credentials.copy()
  1377. if self.system_configuration.credentials
  1378. else {}
  1379. )
  1380. if restrict_model.base_model_name:
  1381. copy_credentials["base_model_name"] = restrict_model.base_model_name
  1382. try:
  1383. custom_model_schema = self.get_model_schema(
  1384. model_type=restrict_model.model_type,
  1385. model=restrict_model.model,
  1386. credentials=copy_credentials,
  1387. )
  1388. except Exception as ex:
  1389. logger.warning("get custom model schema failed, %s", ex)
  1390. continue
  1391. if not custom_model_schema:
  1392. continue
  1393. if custom_model_schema.model_type not in model_types:
  1394. continue
  1395. status = ModelStatus.ACTIVE
  1396. if (
  1397. custom_model_schema.model_type in model_setting_map
  1398. and custom_model_schema.model in model_setting_map[custom_model_schema.model_type]
  1399. ):
  1400. model_setting = model_setting_map[custom_model_schema.model_type][custom_model_schema.model]
  1401. if model_setting.enabled is False:
  1402. status = ModelStatus.DISABLED
  1403. provider_models.append(
  1404. ModelWithProviderEntity(
  1405. model=custom_model_schema.model,
  1406. label=custom_model_schema.label,
  1407. model_type=custom_model_schema.model_type,
  1408. features=custom_model_schema.features,
  1409. fetch_from=FetchFrom.PREDEFINED_MODEL,
  1410. model_properties=custom_model_schema.model_properties,
  1411. deprecated=custom_model_schema.deprecated,
  1412. provider=SimpleModelProviderEntity(self.provider),
  1413. status=status,
  1414. )
  1415. )
  1416. # if llm name not in restricted llm list, remove it
  1417. restrict_model_names = [rm.model for rm in restrict_models]
  1418. for model in provider_models:
  1419. if model.model_type == ModelType.LLM and model.model not in restrict_model_names:
  1420. model.status = ModelStatus.NO_PERMISSION
  1421. elif not quota_configuration.is_valid:
  1422. model.status = ModelStatus.QUOTA_EXCEEDED
  1423. return provider_models
  1424. def _get_custom_provider_models(
  1425. self,
  1426. model_types: Sequence[ModelType],
  1427. provider_schema: ProviderEntity,
  1428. model_setting_map: dict[ModelType, dict[str, ModelSettings]],
  1429. model: str | None = None,
  1430. ) -> list[ModelWithProviderEntity]:
  1431. """
  1432. Get custom provider models.
  1433. :param model_types: model types
  1434. :param provider_schema: provider schema
  1435. :param model_setting_map: model setting map
  1436. :return:
  1437. """
  1438. provider_models = []
  1439. credentials = None
  1440. if self.custom_configuration.provider:
  1441. credentials = self.custom_configuration.provider.credentials
  1442. for model_type in model_types:
  1443. if model_type not in self.provider.supported_model_types:
  1444. continue
  1445. for m in provider_schema.models:
  1446. if m.model_type != model_type:
  1447. continue
  1448. status = ModelStatus.ACTIVE if credentials else ModelStatus.NO_CONFIGURE
  1449. load_balancing_enabled = False
  1450. has_invalid_load_balancing_configs = False
  1451. if m.model_type in model_setting_map and m.model in model_setting_map[m.model_type]:
  1452. model_setting = model_setting_map[m.model_type][m.model]
  1453. if model_setting.enabled is False:
  1454. status = ModelStatus.DISABLED
  1455. provider_model_lb_configs = [
  1456. config
  1457. for config in model_setting.load_balancing_configs
  1458. if config.credential_source_type != "custom_model"
  1459. ]
  1460. load_balancing_enabled = model_setting.load_balancing_enabled
  1461. # when the user enable load_balancing but available configs are less than 2 display warning
  1462. has_invalid_load_balancing_configs = load_balancing_enabled and len(provider_model_lb_configs) < 2
  1463. provider_models.append(
  1464. ModelWithProviderEntity(
  1465. model=m.model,
  1466. label=m.label,
  1467. model_type=m.model_type,
  1468. features=m.features,
  1469. fetch_from=m.fetch_from,
  1470. model_properties=m.model_properties,
  1471. deprecated=m.deprecated,
  1472. provider=SimpleModelProviderEntity(self.provider),
  1473. status=status,
  1474. load_balancing_enabled=load_balancing_enabled,
  1475. has_invalid_load_balancing_configs=has_invalid_load_balancing_configs,
  1476. )
  1477. )
  1478. # custom models
  1479. for model_configuration in self.custom_configuration.models:
  1480. if model_configuration.model_type not in model_types:
  1481. continue
  1482. if model_configuration.unadded_to_model_list:
  1483. continue
  1484. if model and model != model_configuration.model:
  1485. continue
  1486. try:
  1487. custom_model_schema = self.get_model_schema(
  1488. model_type=model_configuration.model_type,
  1489. model=model_configuration.model,
  1490. credentials=model_configuration.credentials,
  1491. )
  1492. except Exception as ex:
  1493. logger.warning("get custom model schema failed, %s", ex)
  1494. continue
  1495. if not custom_model_schema:
  1496. continue
  1497. status = ModelStatus.ACTIVE
  1498. load_balancing_enabled = False
  1499. has_invalid_load_balancing_configs = False
  1500. if (
  1501. custom_model_schema.model_type in model_setting_map
  1502. and custom_model_schema.model in model_setting_map[custom_model_schema.model_type]
  1503. ):
  1504. model_setting = model_setting_map[custom_model_schema.model_type][custom_model_schema.model]
  1505. if model_setting.enabled is False:
  1506. status = ModelStatus.DISABLED
  1507. custom_model_lb_configs = [
  1508. config
  1509. for config in model_setting.load_balancing_configs
  1510. if config.credential_source_type != "provider"
  1511. ]
  1512. load_balancing_enabled = model_setting.load_balancing_enabled
  1513. # when the user enable load_balancing but available configs are less than 2 display warning
  1514. has_invalid_load_balancing_configs = load_balancing_enabled and len(custom_model_lb_configs) < 2
  1515. if len(model_configuration.available_model_credentials) > 0 and not model_configuration.credentials:
  1516. status = ModelStatus.CREDENTIAL_REMOVED
  1517. provider_models.append(
  1518. ModelWithProviderEntity(
  1519. model=custom_model_schema.model,
  1520. label=custom_model_schema.label,
  1521. model_type=custom_model_schema.model_type,
  1522. features=custom_model_schema.features,
  1523. fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,
  1524. model_properties=custom_model_schema.model_properties,
  1525. deprecated=custom_model_schema.deprecated,
  1526. provider=SimpleModelProviderEntity(self.provider),
  1527. status=status,
  1528. load_balancing_enabled=load_balancing_enabled,
  1529. has_invalid_load_balancing_configs=has_invalid_load_balancing_configs,
  1530. )
  1531. )
  1532. return provider_models
  1533. class ProviderConfigurations(BaseModel):
  1534. """
  1535. Model class for provider configuration dict.
  1536. """
  1537. tenant_id: str
  1538. configurations: dict[str, ProviderConfiguration] = Field(default_factory=dict)
  1539. def __init__(self, tenant_id: str):
  1540. super().__init__(tenant_id=tenant_id)
  1541. def get_models(
  1542. self, provider: str | None = None, model_type: ModelType | None = None, only_active: bool = False
  1543. ) -> list[ModelWithProviderEntity]:
  1544. """
  1545. Get available models.
  1546. If preferred provider type is `system`:
  1547. Get the current **system mode** if provider supported,
  1548. if all system modes are not available (no quota), it is considered to be the **custom credential mode**.
  1549. If there is no model configured in custom mode, it is treated as no_configure.
  1550. system > custom > no_configure
  1551. If preferred provider type is `custom`:
  1552. If custom credentials are configured, it is treated as custom mode.
  1553. Otherwise, get the current **system mode** if supported,
  1554. If all system modes are not available (no quota), it is treated as no_configure.
  1555. custom > system > no_configure
  1556. If real mode is `system`, use system credentials to get models,
  1557. paid quotas > provider free quotas > system free quotas
  1558. include pre-defined models (exclude GPT-4, status marked as `no_permission`).
  1559. If real mode is `custom`, use workspace custom credentials to get models,
  1560. include pre-defined models, custom models(manual append).
  1561. If real mode is `no_configure`, only return pre-defined models from `model runtime`.
  1562. (model status marked as `no_configure` if preferred provider type is `custom` otherwise `quota_exceeded`)
  1563. model status marked as `active` is available.
  1564. :param provider: provider name
  1565. :param model_type: model type
  1566. :param only_active: only active models
  1567. :return:
  1568. """
  1569. all_models = []
  1570. for provider_configuration in self.values():
  1571. if provider and provider_configuration.provider.provider != provider:
  1572. continue
  1573. all_models.extend(provider_configuration.get_provider_models(model_type, only_active))
  1574. return all_models
  1575. def to_list(self) -> list[ProviderConfiguration]:
  1576. """
  1577. Convert to list.
  1578. :return:
  1579. """
  1580. return list(self.values())
  1581. def __getitem__(self, key):
  1582. if "/" not in key:
  1583. key = str(ModelProviderID(key))
  1584. return self.configurations[key]
  1585. def __setitem__(self, key, value):
  1586. self.configurations[key] = value
  1587. def __contains__(self, key):
  1588. if "/" not in key:
  1589. key = str(ModelProviderID(key))
  1590. return key in self.configurations
  1591. def __iter__(self):
  1592. # Return an iterator of (key, value) tuples to match BaseModel's __iter__
  1593. yield from self.configurations.items()
  1594. def values(self) -> Iterator[ProviderConfiguration]:
  1595. return iter(self.configurations.values())
  1596. def get(self, key, default=None) -> ProviderConfiguration | None:
  1597. if "/" not in key:
  1598. key = str(ModelProviderID(key))
  1599. return self.configurations.get(key, default) # type: ignore
  1600. class ProviderModelBundle(BaseModel):
  1601. """
  1602. Provider model bundle.
  1603. """
  1604. configuration: ProviderConfiguration
  1605. model_type_instance: AIModel
  1606. # pydantic configs
  1607. model_config = ConfigDict(arbitrary_types_allowed=True, protected_namespaces=())