|
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: |
|
""" |
|
|
|
provider_name_to_provider_records_dict = self._get_all_providers(tenant_id) |
|
|
|
|
|
provider_name_to_provider_records_dict = self._init_trial_provider_records( |
|
tenant_id, provider_name_to_provider_records_dict |
|
) |
|
|
|
|
|
provider_name_to_provider_model_records_dict = self._get_all_provider_models(tenant_id) |
|
|
|
|
|
provider_entities = model_provider_factory.get_providers() |
|
|
|
|
|
provider_name_to_preferred_model_provider_records_dict = self._get_all_preferred_model_providers(tenant_id) |
|
|
|
|
|
provider_name_to_provider_model_settings_dict = self._get_all_provider_model_settings(tenant_id) |
|
|
|
|
|
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) |
|
|
|
|
|
for provider_entity in provider_entities: |
|
|
|
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, []) |
|
|
|
|
|
custom_configuration = self._to_custom_configuration( |
|
tenant_id, provider_entity, provider_records, provider_model_records |
|
) |
|
|
|
|
|
system_configuration = self._to_system_configuration(tenant_id, provider_entity, provider_records) |
|
|
|
|
|
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 |
|
|
|
|
|
provider_model_settings = provider_name_to_provider_model_settings_dict.get(provider_name) |
|
|
|
|
|
provider_load_balancing_configs = provider_name_to_provider_load_balancing_model_configs_dict.get( |
|
provider_name |
|
) |
|
|
|
|
|
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 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) |
|
|
|
|
|
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: |
|
""" |
|
|
|
default_model = ( |
|
db.session.query(TenantDefaultModel) |
|
.filter( |
|
TenantDefaultModel.tenant_id == tenant_id, |
|
TenantDefaultModel.model_type == model_type.to_origin_model_type(), |
|
) |
|
.first() |
|
) |
|
|
|
|
|
|
|
if not default_model: |
|
|
|
provider_configurations = self.get_configurations(tenant_id) |
|
|
|
|
|
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) |
|
|
|
|
|
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.") |
|
|
|
|
|
available_models = provider_configurations.get_models(model_type=model_type, only_active=True) |
|
|
|
|
|
model_names = [model.model for model in available_models] |
|
if model not in model_names: |
|
raise ValueError(f"Model {model} does not exist.") |
|
|
|
|
|
default_model = ( |
|
db.session.query(TenantDefaultModel) |
|
.filter( |
|
TenantDefaultModel.tenant_id == tenant_id, |
|
TenantDefaultModel.model_type == model_type.to_origin_model_type(), |
|
) |
|
.first() |
|
) |
|
|
|
|
|
if default_model: |
|
|
|
default_model.provider_name = provider |
|
default_model.model_name = model |
|
db.session.commit() |
|
else: |
|
|
|
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 |
|
|
|
@staticmethod |
|
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 |
|
|
|
@staticmethod |
|
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: |
|
""" |
|
|
|
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 |
|
|
|
@staticmethod |
|
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 |
|
|
|
@staticmethod |
|
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 |
|
|
|
@staticmethod |
|
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 |
|
|
|
@staticmethod |
|
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: |
|
""" |
|
|
|
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: |
|
|
|
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: |
|
""" |
|
|
|
provider_credential_secret_variables = self._extract_secret_variables( |
|
provider_entity.provider_credential_schema.credential_form_schemas |
|
if provider_entity.provider_credential_schema |
|
else [] |
|
) |
|
|
|
|
|
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 |
|
|
|
|
|
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, |
|
) |
|
|
|
|
|
cached_provider_credentials = provider_credentials_cache.get() |
|
|
|
if not cached_provider_credentials: |
|
try: |
|
|
|
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 = {} |
|
|
|
|
|
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 |
|
|
|
|
|
provider_credentials_cache.set(credentials=provider_credentials) |
|
else: |
|
provider_credentials = cached_provider_credentials |
|
|
|
custom_provider_configuration = CustomProviderConfiguration(credentials=provider_credentials) |
|
|
|
|
|
model_credential_secret_variables = self._extract_secret_variables( |
|
provider_entity.model_credential_schema.credential_form_schemas |
|
if provider_entity.model_credential_schema |
|
else [] |
|
) |
|
|
|
|
|
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 |
|
) |
|
|
|
|
|
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 |
|
|
|
|
|
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 |
|
|
|
|
|
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: |
|
""" |
|
|
|
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) |
|
|
|
|
|
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, |
|
) |
|
|
|
|
|
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 = {} |
|
|
|
|
|
provider_credential_secret_variables = self._extract_secret_variables( |
|
provider_entity.provider_credential_schema.credential_form_schemas |
|
if provider_entity.provider_credential_schema |
|
else [] |
|
) |
|
|
|
|
|
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 |
|
|
|
|
|
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, |
|
) |
|
|
|
@staticmethod |
|
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: |
|
""" |
|
|
|
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") |
|
|
|
@staticmethod |
|
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: |
|
""" |
|
|
|
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, |
|
) |
|
|
|
|
|
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 |
|
|
|
|
|
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 |
|
|
|
|
|
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 |
|
|