Spaces:
Paused
Paused
| import json | |
| from collections import defaultdict | |
| from json import JSONDecodeError | |
| from typing import Optional | |
| from sqlalchemy.exc import IntegrityError | |
| from configs import dify_config | |
| from core.entities.model_entities import DefaultModelEntity, DefaultModelProviderEntity | |
| from core.entities.provider_configuration import ProviderConfiguration, ProviderConfigurations, ProviderModelBundle | |
| from core.entities.provider_entities import ( | |
| CustomConfiguration, | |
| CustomModelConfiguration, | |
| CustomProviderConfiguration, | |
| ModelLoadBalancingConfiguration, | |
| ModelSettings, | |
| QuotaConfiguration, | |
| SystemConfiguration, | |
| ) | |
| from core.helper import encrypter | |
| from core.helper.model_provider_cache import ProviderCredentialsCache, ProviderCredentialsCacheType | |
| from core.helper.position_helper import is_filtered | |
| from core.model_runtime.entities.model_entities import ModelType | |
| from core.model_runtime.entities.provider_entities import CredentialFormSchema, FormType, ProviderEntity | |
| from core.model_runtime.model_providers import model_provider_factory | |
| from extensions import ext_hosting_provider | |
| from extensions.ext_database import db | |
| from extensions.ext_redis import redis_client | |
| from models.provider import ( | |
| LoadBalancingModelConfig, | |
| Provider, | |
| ProviderModel, | |
| ProviderModelSetting, | |
| ProviderQuotaType, | |
| ProviderType, | |
| TenantDefaultModel, | |
| TenantPreferredModelProvider, | |
| ) | |
| from services.feature_service import FeatureService | |
| class ProviderManager: | |
| """ | |
| ProviderManager is a class that manages the model providers includes Hosting and Customize Model Providers. | |
| """ | |
| def __init__(self) -> None: | |
| self.decoding_rsa_key = None | |
| self.decoding_cipher_rsa = None | |
| def get_configurations(self, tenant_id: str) -> ProviderConfigurations: | |
| """ | |
| Get model provider configurations. | |
| Construct ProviderConfiguration objects for each provider | |
| Including: | |
| 1. Basic information of the provider | |
| 2. Hosting configuration information, including: | |
| (1. Whether to enable (support) hosting type, if enabled, the following information exists | |
| (2. List of hosting type provider configurations | |
| (including quota type, quota limit, current remaining quota, etc.) | |
| (3. The current hosting type in use (whether there is a quota or not) | |
| paid quotas > provider free quotas > hosting trial quotas | |
| (4. Unified credentials for hosting providers | |
| 3. Custom configuration information, including: | |
| (1. Whether to enable (support) custom type, if enabled, the following information exists | |
| (2. Custom provider configuration (including credentials) | |
| (3. List of custom provider model configurations (including credentials) | |
| 4. Hosting/custom preferred provider type. | |
| Provide methods: | |
| - Get the current configuration (including credentials) | |
| - Get the availability and status of the hosting configuration: active available, | |
| quota_exceeded insufficient quota, unsupported hosting | |
| - Get the availability of custom configuration | |
| Custom provider available conditions: | |
| (1. custom provider credentials available | |
| (2. at least one custom model credentials available | |
| - Verify, update, and delete custom provider configuration | |
| - Verify, update, and delete custom provider model configuration | |
| - Get the list of available models (optional provider filtering, model type filtering) | |
| Append custom provider models to the list | |
| - Get provider instance | |
| - Switch selection priority | |
| :param tenant_id: | |
| :return: | |
| """ | |
| # Get all provider records of the workspace | |
| provider_name_to_provider_records_dict = self._get_all_providers(tenant_id) | |
| # Initialize trial provider records if not exist | |
| provider_name_to_provider_records_dict = self._init_trial_provider_records( | |
| tenant_id, provider_name_to_provider_records_dict | |
| ) | |
| # Get all provider model records of the workspace | |
| provider_name_to_provider_model_records_dict = self._get_all_provider_models(tenant_id) | |
| # Get all provider entities | |
| provider_entities = model_provider_factory.get_providers() | |
| # Get All preferred provider types of the workspace | |
| provider_name_to_preferred_model_provider_records_dict = self._get_all_preferred_model_providers(tenant_id) | |
| # Get All provider model settings | |
| provider_name_to_provider_model_settings_dict = self._get_all_provider_model_settings(tenant_id) | |
| # Get All load balancing configs | |
| provider_name_to_provider_load_balancing_model_configs_dict = self._get_all_provider_load_balancing_configs( | |
| tenant_id | |
| ) | |
| provider_configurations = ProviderConfigurations(tenant_id=tenant_id) | |
| # Construct ProviderConfiguration objects for each provider | |
| for provider_entity in provider_entities: | |
| # handle include, exclude | |
| if is_filtered( | |
| include_set=dify_config.POSITION_PROVIDER_INCLUDES_SET, | |
| exclude_set=dify_config.POSITION_PROVIDER_EXCLUDES_SET, | |
| data=provider_entity, | |
| name_func=lambda x: x.provider, | |
| ): | |
| continue | |
| provider_name = provider_entity.provider | |
| provider_records = provider_name_to_provider_records_dict.get(provider_entity.provider, []) | |
| provider_model_records = provider_name_to_provider_model_records_dict.get(provider_entity.provider, []) | |
| # Convert to custom configuration | |
| custom_configuration = self._to_custom_configuration( | |
| tenant_id, provider_entity, provider_records, provider_model_records | |
| ) | |
| # Convert to system configuration | |
| system_configuration = self._to_system_configuration(tenant_id, provider_entity, provider_records) | |
| # Get preferred provider type | |
| preferred_provider_type_record = provider_name_to_preferred_model_provider_records_dict.get(provider_name) | |
| if preferred_provider_type_record: | |
| preferred_provider_type = ProviderType.value_of(preferred_provider_type_record.preferred_provider_type) | |
| elif custom_configuration.provider or custom_configuration.models: | |
| preferred_provider_type = ProviderType.CUSTOM | |
| elif system_configuration.enabled: | |
| preferred_provider_type = ProviderType.SYSTEM | |
| else: | |
| preferred_provider_type = ProviderType.CUSTOM | |
| using_provider_type = preferred_provider_type | |
| has_valid_quota = any(quota_conf.is_valid for quota_conf in system_configuration.quota_configurations) | |
| if preferred_provider_type == ProviderType.SYSTEM: | |
| if not system_configuration.enabled or not has_valid_quota: | |
| using_provider_type = ProviderType.CUSTOM | |
| else: | |
| if not custom_configuration.provider and not custom_configuration.models: | |
| if system_configuration.enabled and has_valid_quota: | |
| using_provider_type = ProviderType.SYSTEM | |
| # Get provider load balancing configs | |
| provider_model_settings = provider_name_to_provider_model_settings_dict.get(provider_name) | |
| # Get provider load balancing configs | |
| provider_load_balancing_configs = provider_name_to_provider_load_balancing_model_configs_dict.get( | |
| provider_name | |
| ) | |
| # Convert to model settings | |
| model_settings = self._to_model_settings( | |
| provider_entity=provider_entity, | |
| provider_model_settings=provider_model_settings, | |
| load_balancing_model_configs=provider_load_balancing_configs, | |
| ) | |
| provider_configuration = ProviderConfiguration( | |
| tenant_id=tenant_id, | |
| provider=provider_entity, | |
| preferred_provider_type=preferred_provider_type, | |
| using_provider_type=using_provider_type, | |
| system_configuration=system_configuration, | |
| custom_configuration=custom_configuration, | |
| model_settings=model_settings, | |
| ) | |
| provider_configurations[provider_name] = provider_configuration | |
| # Return the encapsulated object | |
| return provider_configurations | |
| def get_provider_model_bundle(self, tenant_id: str, provider: str, model_type: ModelType) -> ProviderModelBundle: | |
| """ | |
| Get provider model bundle. | |
| :param tenant_id: workspace id | |
| :param provider: provider name | |
| :param model_type: model type | |
| :return: | |
| """ | |
| provider_configurations = self.get_configurations(tenant_id) | |
| # get provider instance | |
| provider_configuration = provider_configurations.get(provider) | |
| if not provider_configuration: | |
| raise ValueError(f"Provider {provider} does not exist.") | |
| provider_instance = provider_configuration.get_provider_instance() | |
| model_type_instance = provider_instance.get_model_instance(model_type) | |
| return ProviderModelBundle( | |
| configuration=provider_configuration, | |
| provider_instance=provider_instance, | |
| model_type_instance=model_type_instance, | |
| ) | |
| def get_default_model(self, tenant_id: str, model_type: ModelType) -> Optional[DefaultModelEntity]: | |
| """ | |
| Get default model. | |
| :param tenant_id: workspace id | |
| :param model_type: model type | |
| :return: | |
| """ | |
| # Get the corresponding TenantDefaultModel record | |
| default_model = ( | |
| db.session.query(TenantDefaultModel) | |
| .filter( | |
| TenantDefaultModel.tenant_id == tenant_id, | |
| TenantDefaultModel.model_type == model_type.to_origin_model_type(), | |
| ) | |
| .first() | |
| ) | |
| # If it does not exist, get the first available provider model from get_configurations | |
| # and update the TenantDefaultModel record | |
| if not default_model: | |
| # Get provider configurations | |
| provider_configurations = self.get_configurations(tenant_id) | |
| # get available models from provider_configurations | |
| available_models = provider_configurations.get_models(model_type=model_type, only_active=True) | |
| if available_models: | |
| available_model = next( | |
| (model for model in available_models if model.model == "gpt-4"), available_models[0] | |
| ) | |
| default_model = TenantDefaultModel( | |
| tenant_id=tenant_id, | |
| model_type=model_type.to_origin_model_type(), | |
| provider_name=available_model.provider.provider, | |
| model_name=available_model.model, | |
| ) | |
| db.session.add(default_model) | |
| db.session.commit() | |
| if not default_model: | |
| return None | |
| provider_instance = model_provider_factory.get_provider_instance(default_model.provider_name) | |
| provider_schema = provider_instance.get_provider_schema() | |
| return DefaultModelEntity( | |
| model=default_model.model_name, | |
| model_type=model_type, | |
| provider=DefaultModelProviderEntity( | |
| provider=provider_schema.provider, | |
| label=provider_schema.label, | |
| icon_small=provider_schema.icon_small, | |
| icon_large=provider_schema.icon_large, | |
| supported_model_types=provider_schema.supported_model_types, | |
| ), | |
| ) | |
| def get_first_provider_first_model(self, tenant_id: str, model_type: ModelType) -> tuple[str, str]: | |
| """ | |
| Get names of first model and its provider | |
| :param tenant_id: workspace id | |
| :param model_type: model type | |
| :return: provider name, model name | |
| """ | |
| provider_configurations = self.get_configurations(tenant_id) | |
| # get available models from provider_configurations | |
| all_models = provider_configurations.get_models(model_type=model_type, only_active=False) | |
| return all_models[0].provider.provider, all_models[0].model | |
| def update_default_model_record( | |
| self, tenant_id: str, model_type: ModelType, provider: str, model: str | |
| ) -> TenantDefaultModel: | |
| """ | |
| Update default model record. | |
| :param tenant_id: workspace id | |
| :param model_type: model type | |
| :param provider: provider name | |
| :param model: model name | |
| :return: | |
| """ | |
| provider_configurations = self.get_configurations(tenant_id) | |
| if provider not in provider_configurations: | |
| raise ValueError(f"Provider {provider} does not exist.") | |
| # get available models from provider_configurations | |
| available_models = provider_configurations.get_models(model_type=model_type, only_active=True) | |
| # check if the model is exist in available models | |
| model_names = [model.model for model in available_models] | |
| if model not in model_names: | |
| raise ValueError(f"Model {model} does not exist.") | |
| # Get the list of available models from get_configurations and check if it is LLM | |
| default_model = ( | |
| db.session.query(TenantDefaultModel) | |
| .filter( | |
| TenantDefaultModel.tenant_id == tenant_id, | |
| TenantDefaultModel.model_type == model_type.to_origin_model_type(), | |
| ) | |
| .first() | |
| ) | |
| # create or update TenantDefaultModel record | |
| if default_model: | |
| # update default model | |
| default_model.provider_name = provider | |
| default_model.model_name = model | |
| db.session.commit() | |
| else: | |
| # create default model | |
| default_model = TenantDefaultModel( | |
| tenant_id=tenant_id, | |
| model_type=model_type.value, | |
| provider_name=provider, | |
| model_name=model, | |
| ) | |
| db.session.add(default_model) | |
| db.session.commit() | |
| return default_model | |
| def _get_all_providers(tenant_id: str) -> dict[str, list[Provider]]: | |
| """ | |
| Get all provider records of the workspace. | |
| :param tenant_id: workspace id | |
| :return: | |
| """ | |
| providers = db.session.query(Provider).filter(Provider.tenant_id == tenant_id, Provider.is_valid == True).all() | |
| provider_name_to_provider_records_dict = defaultdict(list) | |
| for provider in providers: | |
| provider_name_to_provider_records_dict[provider.provider_name].append(provider) | |
| return provider_name_to_provider_records_dict | |
| def _get_all_provider_models(tenant_id: str) -> dict[str, list[ProviderModel]]: | |
| """ | |
| Get all provider model records of the workspace. | |
| :param tenant_id: workspace id | |
| :return: | |
| """ | |
| # Get all provider model records of the workspace | |
| provider_models = ( | |
| db.session.query(ProviderModel) | |
| .filter(ProviderModel.tenant_id == tenant_id, ProviderModel.is_valid == True) | |
| .all() | |
| ) | |
| provider_name_to_provider_model_records_dict = defaultdict(list) | |
| for provider_model in provider_models: | |
| provider_name_to_provider_model_records_dict[provider_model.provider_name].append(provider_model) | |
| return provider_name_to_provider_model_records_dict | |
| def _get_all_preferred_model_providers(tenant_id: str) -> dict[str, TenantPreferredModelProvider]: | |
| """ | |
| Get All preferred provider types of the workspace. | |
| :param tenant_id: workspace id | |
| :return: | |
| """ | |
| preferred_provider_types = ( | |
| db.session.query(TenantPreferredModelProvider) | |
| .filter(TenantPreferredModelProvider.tenant_id == tenant_id) | |
| .all() | |
| ) | |
| provider_name_to_preferred_provider_type_records_dict = { | |
| preferred_provider_type.provider_name: preferred_provider_type | |
| for preferred_provider_type in preferred_provider_types | |
| } | |
| return provider_name_to_preferred_provider_type_records_dict | |
| def _get_all_provider_model_settings(tenant_id: str) -> dict[str, list[ProviderModelSetting]]: | |
| """ | |
| Get All provider model settings of the workspace. | |
| :param tenant_id: workspace id | |
| :return: | |
| """ | |
| provider_model_settings = ( | |
| db.session.query(ProviderModelSetting).filter(ProviderModelSetting.tenant_id == tenant_id).all() | |
| ) | |
| provider_name_to_provider_model_settings_dict = defaultdict(list) | |
| for provider_model_setting in provider_model_settings: | |
| ( | |
| provider_name_to_provider_model_settings_dict[provider_model_setting.provider_name].append( | |
| provider_model_setting | |
| ) | |
| ) | |
| return provider_name_to_provider_model_settings_dict | |
| def _get_all_provider_load_balancing_configs(tenant_id: str) -> dict[str, list[LoadBalancingModelConfig]]: | |
| """ | |
| Get All provider load balancing configs of the workspace. | |
| :param tenant_id: workspace id | |
| :return: | |
| """ | |
| cache_key = f"tenant:{tenant_id}:model_load_balancing_enabled" | |
| cache_result = redis_client.get(cache_key) | |
| if cache_result is None: | |
| model_load_balancing_enabled = FeatureService.get_features(tenant_id).model_load_balancing_enabled | |
| redis_client.setex(cache_key, 120, str(model_load_balancing_enabled)) | |
| else: | |
| cache_result = cache_result.decode("utf-8") | |
| model_load_balancing_enabled = cache_result == "True" | |
| if not model_load_balancing_enabled: | |
| return {} | |
| provider_load_balancing_configs = ( | |
| db.session.query(LoadBalancingModelConfig).filter(LoadBalancingModelConfig.tenant_id == tenant_id).all() | |
| ) | |
| provider_name_to_provider_load_balancing_model_configs_dict = defaultdict(list) | |
| for provider_load_balancing_config in provider_load_balancing_configs: | |
| ( | |
| provider_name_to_provider_load_balancing_model_configs_dict[ | |
| provider_load_balancing_config.provider_name | |
| ].append(provider_load_balancing_config) | |
| ) | |
| return provider_name_to_provider_load_balancing_model_configs_dict | |
| def _init_trial_provider_records( | |
| tenant_id: str, provider_name_to_provider_records_dict: dict[str, list] | |
| ) -> dict[str, list]: | |
| """ | |
| Initialize trial provider records if not exists. | |
| :param tenant_id: workspace id | |
| :param provider_name_to_provider_records_dict: provider name to provider records dict | |
| :return: | |
| """ | |
| # Get hosting configuration | |
| hosting_configuration = ext_hosting_provider.hosting_configuration | |
| for provider_name, configuration in hosting_configuration.provider_map.items(): | |
| if not configuration.enabled: | |
| continue | |
| provider_records = provider_name_to_provider_records_dict.get(provider_name) | |
| if not provider_records: | |
| provider_records = [] | |
| provider_quota_to_provider_record_dict = {} | |
| for provider_record in provider_records: | |
| if provider_record.provider_type != ProviderType.SYSTEM.value: | |
| continue | |
| provider_quota_to_provider_record_dict[ProviderQuotaType.value_of(provider_record.quota_type)] = ( | |
| provider_record | |
| ) | |
| for quota in configuration.quotas: | |
| if quota.quota_type == ProviderQuotaType.TRIAL: | |
| # Init trial provider records if not exists | |
| if ProviderQuotaType.TRIAL not in provider_quota_to_provider_record_dict: | |
| try: | |
| provider_record = Provider( | |
| tenant_id=tenant_id, | |
| provider_name=provider_name, | |
| provider_type=ProviderType.SYSTEM.value, | |
| quota_type=ProviderQuotaType.TRIAL.value, | |
| quota_limit=quota.quota_limit, | |
| quota_used=0, | |
| is_valid=True, | |
| ) | |
| db.session.add(provider_record) | |
| db.session.commit() | |
| except IntegrityError: | |
| db.session.rollback() | |
| provider_record = ( | |
| db.session.query(Provider) | |
| .filter( | |
| Provider.tenant_id == tenant_id, | |
| Provider.provider_name == provider_name, | |
| Provider.provider_type == ProviderType.SYSTEM.value, | |
| Provider.quota_type == ProviderQuotaType.TRIAL.value, | |
| ) | |
| .first() | |
| ) | |
| if provider_record and not provider_record.is_valid: | |
| provider_record.is_valid = True | |
| db.session.commit() | |
| provider_name_to_provider_records_dict[provider_name].append(provider_record) | |
| return provider_name_to_provider_records_dict | |
| def _to_custom_configuration( | |
| self, | |
| tenant_id: str, | |
| provider_entity: ProviderEntity, | |
| provider_records: list[Provider], | |
| provider_model_records: list[ProviderModel], | |
| ) -> CustomConfiguration: | |
| """ | |
| Convert to custom configuration. | |
| :param tenant_id: workspace id | |
| :param provider_entity: provider entity | |
| :param provider_records: provider records | |
| :param provider_model_records: provider model records | |
| :return: | |
| """ | |
| # Get provider credential secret variables | |
| provider_credential_secret_variables = self._extract_secret_variables( | |
| provider_entity.provider_credential_schema.credential_form_schemas | |
| if provider_entity.provider_credential_schema | |
| else [] | |
| ) | |
| # Get custom provider record | |
| custom_provider_record = None | |
| for provider_record in provider_records: | |
| if provider_record.provider_type == ProviderType.SYSTEM.value: | |
| continue | |
| if not provider_record.encrypted_config: | |
| continue | |
| custom_provider_record = provider_record | |
| # Get custom provider credentials | |
| custom_provider_configuration = None | |
| if custom_provider_record: | |
| provider_credentials_cache = ProviderCredentialsCache( | |
| tenant_id=tenant_id, | |
| identity_id=custom_provider_record.id, | |
| cache_type=ProviderCredentialsCacheType.PROVIDER, | |
| ) | |
| # Get cached provider credentials | |
| cached_provider_credentials = provider_credentials_cache.get() | |
| if not cached_provider_credentials: | |
| try: | |
| # fix origin data | |
| if ( | |
| custom_provider_record.encrypted_config | |
| and not custom_provider_record.encrypted_config.startswith("{") | |
| ): | |
| provider_credentials = {"openai_api_key": custom_provider_record.encrypted_config} | |
| else: | |
| provider_credentials = json.loads(custom_provider_record.encrypted_config) | |
| except JSONDecodeError: | |
| provider_credentials = {} | |
| # Get decoding rsa key and cipher for decrypting credentials | |
| if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None: | |
| self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id) | |
| for variable in provider_credential_secret_variables: | |
| if variable in provider_credentials: | |
| try: | |
| provider_credentials[variable] = encrypter.decrypt_token_with_decoding( | |
| provider_credentials.get(variable), self.decoding_rsa_key, self.decoding_cipher_rsa | |
| ) | |
| except ValueError: | |
| pass | |
| # cache provider credentials | |
| provider_credentials_cache.set(credentials=provider_credentials) | |
| else: | |
| provider_credentials = cached_provider_credentials | |
| custom_provider_configuration = CustomProviderConfiguration(credentials=provider_credentials) | |
| # Get provider model credential secret variables | |
| model_credential_secret_variables = self._extract_secret_variables( | |
| provider_entity.model_credential_schema.credential_form_schemas | |
| if provider_entity.model_credential_schema | |
| else [] | |
| ) | |
| # Get custom provider model credentials | |
| custom_model_configurations = [] | |
| for provider_model_record in provider_model_records: | |
| if not provider_model_record.encrypted_config: | |
| continue | |
| provider_model_credentials_cache = ProviderCredentialsCache( | |
| tenant_id=tenant_id, identity_id=provider_model_record.id, cache_type=ProviderCredentialsCacheType.MODEL | |
| ) | |
| # Get cached provider model credentials | |
| cached_provider_model_credentials = provider_model_credentials_cache.get() | |
| if not cached_provider_model_credentials: | |
| try: | |
| provider_model_credentials = json.loads(provider_model_record.encrypted_config) | |
| except JSONDecodeError: | |
| continue | |
| # Get decoding rsa key and cipher for decrypting credentials | |
| if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None: | |
| self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id) | |
| for variable in model_credential_secret_variables: | |
| if variable in provider_model_credentials: | |
| try: | |
| provider_model_credentials[variable] = encrypter.decrypt_token_with_decoding( | |
| provider_model_credentials.get(variable), | |
| self.decoding_rsa_key, | |
| self.decoding_cipher_rsa, | |
| ) | |
| except ValueError: | |
| pass | |
| # cache provider model credentials | |
| provider_model_credentials_cache.set(credentials=provider_model_credentials) | |
| else: | |
| provider_model_credentials = cached_provider_model_credentials | |
| custom_model_configurations.append( | |
| CustomModelConfiguration( | |
| model=provider_model_record.model_name, | |
| model_type=ModelType.value_of(provider_model_record.model_type), | |
| credentials=provider_model_credentials, | |
| ) | |
| ) | |
| return CustomConfiguration(provider=custom_provider_configuration, models=custom_model_configurations) | |
| def _to_system_configuration( | |
| self, tenant_id: str, provider_entity: ProviderEntity, provider_records: list[Provider] | |
| ) -> SystemConfiguration: | |
| """ | |
| Convert to system configuration. | |
| :param tenant_id: workspace id | |
| :param provider_entity: provider entity | |
| :param provider_records: provider records | |
| :return: | |
| """ | |
| # Get hosting configuration | |
| hosting_configuration = ext_hosting_provider.hosting_configuration | |
| if ( | |
| provider_entity.provider not in hosting_configuration.provider_map | |
| or not hosting_configuration.provider_map.get(provider_entity.provider).enabled | |
| ): | |
| return SystemConfiguration(enabled=False) | |
| provider_hosting_configuration = hosting_configuration.provider_map.get(provider_entity.provider) | |
| # Convert provider_records to dict | |
| quota_type_to_provider_records_dict = {} | |
| for provider_record in provider_records: | |
| if provider_record.provider_type != ProviderType.SYSTEM.value: | |
| continue | |
| quota_type_to_provider_records_dict[ProviderQuotaType.value_of(provider_record.quota_type)] = ( | |
| provider_record | |
| ) | |
| quota_configurations = [] | |
| for provider_quota in provider_hosting_configuration.quotas: | |
| if provider_quota.quota_type not in quota_type_to_provider_records_dict: | |
| if provider_quota.quota_type == ProviderQuotaType.FREE: | |
| quota_configuration = QuotaConfiguration( | |
| quota_type=provider_quota.quota_type, | |
| quota_unit=provider_hosting_configuration.quota_unit, | |
| quota_used=0, | |
| quota_limit=0, | |
| is_valid=False, | |
| restrict_models=provider_quota.restrict_models, | |
| ) | |
| else: | |
| continue | |
| else: | |
| provider_record = quota_type_to_provider_records_dict[provider_quota.quota_type] | |
| quota_configuration = QuotaConfiguration( | |
| quota_type=provider_quota.quota_type, | |
| quota_unit=provider_hosting_configuration.quota_unit, | |
| quota_used=provider_record.quota_used, | |
| quota_limit=provider_record.quota_limit, | |
| is_valid=provider_record.quota_limit > provider_record.quota_used | |
| or provider_record.quota_limit == -1, | |
| restrict_models=provider_quota.restrict_models, | |
| ) | |
| quota_configurations.append(quota_configuration) | |
| if len(quota_configurations) == 0: | |
| return SystemConfiguration(enabled=False) | |
| current_quota_type = self._choice_current_using_quota_type(quota_configurations) | |
| current_using_credentials = provider_hosting_configuration.credentials | |
| if current_quota_type == ProviderQuotaType.FREE: | |
| provider_record = quota_type_to_provider_records_dict.get(current_quota_type) | |
| if provider_record: | |
| provider_credentials_cache = ProviderCredentialsCache( | |
| tenant_id=tenant_id, | |
| identity_id=provider_record.id, | |
| cache_type=ProviderCredentialsCacheType.PROVIDER, | |
| ) | |
| # Get cached provider credentials | |
| cached_provider_credentials = provider_credentials_cache.get() | |
| if not cached_provider_credentials: | |
| try: | |
| provider_credentials = json.loads(provider_record.encrypted_config) | |
| except JSONDecodeError: | |
| provider_credentials = {} | |
| # Get provider credential secret variables | |
| provider_credential_secret_variables = self._extract_secret_variables( | |
| provider_entity.provider_credential_schema.credential_form_schemas | |
| if provider_entity.provider_credential_schema | |
| else [] | |
| ) | |
| # Get decoding rsa key and cipher for decrypting credentials | |
| if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None: | |
| self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id) | |
| for variable in provider_credential_secret_variables: | |
| if variable in provider_credentials: | |
| try: | |
| provider_credentials[variable] = encrypter.decrypt_token_with_decoding( | |
| provider_credentials.get(variable), self.decoding_rsa_key, self.decoding_cipher_rsa | |
| ) | |
| except ValueError: | |
| pass | |
| current_using_credentials = provider_credentials | |
| # cache provider credentials | |
| provider_credentials_cache.set(credentials=current_using_credentials) | |
| else: | |
| current_using_credentials = cached_provider_credentials | |
| else: | |
| current_using_credentials = {} | |
| quota_configurations = [] | |
| return SystemConfiguration( | |
| enabled=True, | |
| current_quota_type=current_quota_type, | |
| quota_configurations=quota_configurations, | |
| credentials=current_using_credentials, | |
| ) | |
| def _choice_current_using_quota_type(quota_configurations: list[QuotaConfiguration]) -> ProviderQuotaType: | |
| """ | |
| Choice current using quota type. | |
| paid quotas > provider free quotas > hosting trial quotas | |
| If there is still quota for the corresponding quota type according to the sorting, | |
| :param quota_configurations: | |
| :return: | |
| """ | |
| # convert to dict | |
| quota_type_to_quota_configuration_dict = { | |
| quota_configuration.quota_type: quota_configuration for quota_configuration in quota_configurations | |
| } | |
| last_quota_configuration = None | |
| for quota_type in [ProviderQuotaType.PAID, ProviderQuotaType.FREE, ProviderQuotaType.TRIAL]: | |
| if quota_type in quota_type_to_quota_configuration_dict: | |
| last_quota_configuration = quota_type_to_quota_configuration_dict[quota_type] | |
| if last_quota_configuration.is_valid: | |
| return quota_type | |
| if last_quota_configuration: | |
| return last_quota_configuration.quota_type | |
| raise ValueError("No quota type available") | |
| def _extract_secret_variables(credential_form_schemas: list[CredentialFormSchema]) -> list[str]: | |
| """ | |
| Extract secret input form variables. | |
| :param credential_form_schemas: | |
| :return: | |
| """ | |
| secret_input_form_variables = [] | |
| for credential_form_schema in credential_form_schemas: | |
| if credential_form_schema.type == FormType.SECRET_INPUT: | |
| secret_input_form_variables.append(credential_form_schema.variable) | |
| return secret_input_form_variables | |
| def _to_model_settings( | |
| self, | |
| provider_entity: ProviderEntity, | |
| provider_model_settings: Optional[list[ProviderModelSetting]] = None, | |
| load_balancing_model_configs: Optional[list[LoadBalancingModelConfig]] = None, | |
| ) -> list[ModelSettings]: | |
| """ | |
| Convert to model settings. | |
| :param provider_entity: provider entity | |
| :param provider_model_settings: provider model settings include enabled, load balancing enabled | |
| :param load_balancing_model_configs: load balancing model configs | |
| :return: | |
| """ | |
| # Get provider model credential secret variables | |
| model_credential_secret_variables = self._extract_secret_variables( | |
| provider_entity.model_credential_schema.credential_form_schemas | |
| if provider_entity.model_credential_schema | |
| else [] | |
| ) | |
| model_settings = [] | |
| if not provider_model_settings: | |
| return model_settings | |
| for provider_model_setting in provider_model_settings: | |
| load_balancing_configs = [] | |
| if provider_model_setting.load_balancing_enabled and load_balancing_model_configs: | |
| for load_balancing_model_config in load_balancing_model_configs: | |
| if ( | |
| load_balancing_model_config.model_name == provider_model_setting.model_name | |
| and load_balancing_model_config.model_type == provider_model_setting.model_type | |
| ): | |
| if not load_balancing_model_config.enabled: | |
| continue | |
| if not load_balancing_model_config.encrypted_config: | |
| if load_balancing_model_config.name == "__inherit__": | |
| load_balancing_configs.append( | |
| ModelLoadBalancingConfiguration( | |
| id=load_balancing_model_config.id, | |
| name=load_balancing_model_config.name, | |
| credentials={}, | |
| ) | |
| ) | |
| continue | |
| provider_model_credentials_cache = ProviderCredentialsCache( | |
| tenant_id=load_balancing_model_config.tenant_id, | |
| identity_id=load_balancing_model_config.id, | |
| cache_type=ProviderCredentialsCacheType.LOAD_BALANCING_MODEL, | |
| ) | |
| # Get cached provider model credentials | |
| cached_provider_model_credentials = provider_model_credentials_cache.get() | |
| if not cached_provider_model_credentials: | |
| try: | |
| provider_model_credentials = json.loads(load_balancing_model_config.encrypted_config) | |
| except JSONDecodeError: | |
| continue | |
| # Get decoding rsa key and cipher for decrypting credentials | |
| if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None: | |
| self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding( | |
| load_balancing_model_config.tenant_id | |
| ) | |
| for variable in model_credential_secret_variables: | |
| if variable in provider_model_credentials: | |
| try: | |
| provider_model_credentials[variable] = encrypter.decrypt_token_with_decoding( | |
| provider_model_credentials.get(variable), | |
| self.decoding_rsa_key, | |
| self.decoding_cipher_rsa, | |
| ) | |
| except ValueError: | |
| pass | |
| # cache provider model credentials | |
| provider_model_credentials_cache.set(credentials=provider_model_credentials) | |
| else: | |
| provider_model_credentials = cached_provider_model_credentials | |
| load_balancing_configs.append( | |
| ModelLoadBalancingConfiguration( | |
| id=load_balancing_model_config.id, | |
| name=load_balancing_model_config.name, | |
| credentials=provider_model_credentials, | |
| ) | |
| ) | |
| model_settings.append( | |
| ModelSettings( | |
| model=provider_model_setting.model_name, | |
| model_type=ModelType.value_of(provider_model_setting.model_type), | |
| enabled=provider_model_setting.enabled, | |
| load_balancing_configs=load_balancing_configs if len(load_balancing_configs) > 1 else [], | |
| ) | |
| ) | |
| return model_settings | |