dify

dify/api/services/model_load_balancing_service.py

@@ -0,0 +1,620 @@
+import json
+import logging
+from json import JSONDecodeError
+from typing import Union
+
+from sqlalchemy import or_, select
+
+from constants import HIDDEN_VALUE
+from core.entities.provider_configuration import ProviderConfiguration
+from core.helper import encrypter
+from core.helper.model_provider_cache import ProviderCredentialsCache, ProviderCredentialsCacheType
+from core.model_manager import LBModelManager
+from core.model_runtime.entities.model_entities import ModelType
+from core.model_runtime.entities.provider_entities import (
+    ModelCredentialSchema,
+    ProviderCredentialSchema,
+)
+from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
+from core.provider_manager import ProviderManager
+from extensions.ext_database import db
+from libs.datetime_utils import naive_utc_now
+from models.provider import LoadBalancingModelConfig, ProviderCredential, ProviderModelCredential
+
+logger = logging.getLogger(__name__)
+
+
+class ModelLoadBalancingService:
+    def __init__(self):
+        self.provider_manager = ProviderManager()
+
+    def enable_model_load_balancing(self, tenant_id: str, provider: str, model: str, model_type: str):
+        """
+        enable model load balancing.
+
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model: model name
+        :param model_type: model type
+        :return:
+        """
+        # Get all provider configurations of the current workspace
+        provider_configurations = self.provider_manager.get_configurations(tenant_id)
+
+        # Get provider configuration
+        provider_configuration = provider_configurations.get(provider)
+        if not provider_configuration:
+            raise ValueError(f"Provider {provider} does not exist.")
+
+        # Enable model load balancing
+        provider_configuration.enable_model_load_balancing(model=model, model_type=ModelType.value_of(model_type))
+
+    def disable_model_load_balancing(self, tenant_id: str, provider: str, model: str, model_type: str):
+        """
+        disable model load balancing.
+
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model: model name
+        :param model_type: model type
+        :return:
+        """
+        # Get all provider configurations of the current workspace
+        provider_configurations = self.provider_manager.get_configurations(tenant_id)
+
+        # Get provider configuration
+        provider_configuration = provider_configurations.get(provider)
+        if not provider_configuration:
+            raise ValueError(f"Provider {provider} does not exist.")
+
+        # disable model load balancing
+        provider_configuration.disable_model_load_balancing(model=model, model_type=ModelType.value_of(model_type))
+
+    def get_load_balancing_configs(
+        self, tenant_id: str, provider: str, model: str, model_type: str, config_from: str = ""
+    ) -> tuple[bool, list[dict]]:
+        """
+        Get load balancing configurations.
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model: model name
+        :param model_type: model type
+        :return:
+        """
+        # Get all provider configurations of the current workspace
+        provider_configurations = self.provider_manager.get_configurations(tenant_id)
+
+        # Get provider configuration
+        provider_configuration = provider_configurations.get(provider)
+        if not provider_configuration:
+            raise ValueError(f"Provider {provider} does not exist.")
+
+        # Convert model type to ModelType
+        model_type_enum = ModelType.value_of(model_type)
+
+        # Get provider model setting
+        provider_model_setting = provider_configuration.get_provider_model_setting(
+            model_type=model_type_enum,
+            model=model,
+        )
+
+        is_load_balancing_enabled = False
+        if provider_model_setting and provider_model_setting.load_balancing_enabled:
+            is_load_balancing_enabled = True
+
+        if config_from == "predefined-model":
+            credential_source_type = "provider"
+        else:
+            credential_source_type = "custom_model"
+
+        # Get load balancing configurations
+        load_balancing_configs = (
+            db.session.query(LoadBalancingModelConfig)
+            .where(
+                LoadBalancingModelConfig.tenant_id == tenant_id,
+                LoadBalancingModelConfig.provider_name == provider_configuration.provider.provider,
+                LoadBalancingModelConfig.model_type == model_type_enum.to_origin_model_type(),
+                LoadBalancingModelConfig.model_name == model,
+                or_(
+                    LoadBalancingModelConfig.credential_source_type == credential_source_type,
+                    LoadBalancingModelConfig.credential_source_type.is_(None),
+                ),
+            )
+            .order_by(LoadBalancingModelConfig.created_at)
+            .all()
+        )
+
+        if provider_configuration.custom_configuration.provider:
+            # check if the inherit configuration exists,
+            # inherit is represented for the provider or model custom credentials
+            inherit_config_exists = False
+            for load_balancing_config in load_balancing_configs:
+                if load_balancing_config.name == "__inherit__":
+                    inherit_config_exists = True
+                    break
+
+            if not inherit_config_exists:
+                # Initialize the inherit configuration
+                inherit_config = self._init_inherit_config(tenant_id, provider, model, model_type_enum)
+
+                # prepend the inherit configuration
+                load_balancing_configs.insert(0, inherit_config)
+            else:
+                # move the inherit configuration to the first
+                for i, load_balancing_config in enumerate(load_balancing_configs[:]):
+                    if load_balancing_config.name == "__inherit__":
+                        inherit_config = load_balancing_configs.pop(i)
+                        load_balancing_configs.insert(0, inherit_config)
+
+        # Get credential form schemas from model credential schema or provider credential schema
+        credential_schemas = self._get_credential_schema(provider_configuration)
+
+        # Get decoding rsa key and cipher for decrypting credentials
+        decoding_rsa_key, decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id)
+
+        # fetch status and ttl for each config
+        datas = []
+        for load_balancing_config in load_balancing_configs:
+            in_cooldown, ttl = LBModelManager.get_config_in_cooldown_and_ttl(
+                tenant_id=tenant_id,
+                provider=provider,
+                model=model,
+                model_type=model_type_enum,
+                config_id=load_balancing_config.id,
+            )
+
+            try:
+                if load_balancing_config.encrypted_config:
+                    credentials: dict[str, object] = json.loads(load_balancing_config.encrypted_config)
+                else:
+                    credentials = {}
+            except JSONDecodeError:
+                credentials = {}
+
+            # Get provider credential secret variables
+            credential_secret_variables = provider_configuration.extract_secret_variables(
+                credential_schemas.credential_form_schemas
+            )
+
+            # decrypt credentials
+            for variable in credential_secret_variables:
+                if variable in credentials:
+                    try:
+                        token_value = credentials.get(variable)
+                        if isinstance(token_value, str):
+                            credentials[variable] = encrypter.decrypt_token_with_decoding(
+                                token_value,
+                                decoding_rsa_key,
+                                decoding_cipher_rsa,
+                            )
+                    except ValueError:
+                        pass
+
+            # Obfuscate credentials
+            credentials = provider_configuration.obfuscated_credentials(
+                credentials=credentials, credential_form_schemas=credential_schemas.credential_form_schemas
+            )
+
+            datas.append(
+                {
+                    "id": load_balancing_config.id,
+                    "name": load_balancing_config.name,
+                    "credentials": credentials,
+                    "credential_id": load_balancing_config.credential_id,
+                    "enabled": load_balancing_config.enabled,
+                    "in_cooldown": in_cooldown,
+                    "ttl": ttl,
+                }
+            )
+
+        return is_load_balancing_enabled, datas
+
+    def get_load_balancing_config(
+        self, tenant_id: str, provider: str, model: str, model_type: str, config_id: str
+    ) -> dict | None:
+        """
+        Get load balancing configuration.
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model: model name
+        :param model_type: model type
+        :param config_id: load balancing config id
+        :return:
+        """
+        # Get all provider configurations of the current workspace
+        provider_configurations = self.provider_manager.get_configurations(tenant_id)
+
+        # Get provider configuration
+        provider_configuration = provider_configurations.get(provider)
+        if not provider_configuration:
+            raise ValueError(f"Provider {provider} does not exist.")
+
+        # Convert model type to ModelType
+        model_type_enum = ModelType.value_of(model_type)
+
+        # Get load balancing configurations
+        load_balancing_model_config = (
+            db.session.query(LoadBalancingModelConfig)
+            .where(
+                LoadBalancingModelConfig.tenant_id == tenant_id,
+                LoadBalancingModelConfig.provider_name == provider_configuration.provider.provider,
+                LoadBalancingModelConfig.model_type == model_type_enum.to_origin_model_type(),
+                LoadBalancingModelConfig.model_name == model,
+                LoadBalancingModelConfig.id == config_id,
+            )
+            .first()
+        )
+
+        if not load_balancing_model_config:
+            return None
+
+        try:
+            if load_balancing_model_config.encrypted_config:
+                credentials = json.loads(load_balancing_model_config.encrypted_config)
+            else:
+                credentials = {}
+        except JSONDecodeError:
+            credentials = {}
+
+        # Get credential form schemas from model credential schema or provider credential schema
+        credential_schemas = self._get_credential_schema(provider_configuration)
+
+        # Obfuscate credentials
+        credentials = provider_configuration.obfuscated_credentials(
+            credentials=credentials, credential_form_schemas=credential_schemas.credential_form_schemas
+        )
+
+        return {
+            "id": load_balancing_model_config.id,
+            "name": load_balancing_model_config.name,
+            "credentials": credentials,
+            "enabled": load_balancing_model_config.enabled,
+        }
+
+    def _init_inherit_config(
+        self, tenant_id: str, provider: str, model: str, model_type: ModelType
+    ) -> LoadBalancingModelConfig:
+        """
+        Initialize the inherit configuration.
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model: model name
+        :param model_type: model type
+        :return:
+        """
+        # Initialize the inherit configuration
+        inherit_config = LoadBalancingModelConfig(
+            tenant_id=tenant_id,
+            provider_name=provider,
+            model_type=model_type.to_origin_model_type(),
+            model_name=model,
+            name="__inherit__",
+        )
+        db.session.add(inherit_config)
+        db.session.commit()
+
+        return inherit_config
+
+    def update_load_balancing_configs(
+        self, tenant_id: str, provider: str, model: str, model_type: str, configs: list[dict], config_from: str
+    ):
+        """
+        Update load balancing configurations.
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model: model name
+        :param model_type: model type
+        :param configs: load balancing configs
+        :param config_from: predefined-model or custom-model
+        :return:
+        """
+        # Get all provider configurations of the current workspace
+        provider_configurations = self.provider_manager.get_configurations(tenant_id)
+
+        # Get provider configuration
+        provider_configuration = provider_configurations.get(provider)
+        if not provider_configuration:
+            raise ValueError(f"Provider {provider} does not exist.")
+
+        # Convert model type to ModelType
+        model_type_enum = ModelType.value_of(model_type)
+
+        if not isinstance(configs, list):
+            raise ValueError("Invalid load balancing configs")
+
+        current_load_balancing_configs = db.session.scalars(
+            select(LoadBalancingModelConfig).where(
+                LoadBalancingModelConfig.tenant_id == tenant_id,
+                LoadBalancingModelConfig.provider_name == provider_configuration.provider.provider,
+                LoadBalancingModelConfig.model_type == model_type_enum.to_origin_model_type(),
+                LoadBalancingModelConfig.model_name == model,
+            )
+        ).all()
+
+        # id as key, config as value
+        current_load_balancing_configs_dict = {config.id: config for config in current_load_balancing_configs}
+        updated_config_ids = set()
+
+        for config in configs:
+            if not isinstance(config, dict):
+                raise ValueError("Invalid load balancing config")
+
+            config_id = config.get("id")
+            name = config.get("name")
+            credentials = config.get("credentials")
+            credential_id = config.get("credential_id")
+            enabled = config.get("enabled")
+
+            credential_record: ProviderCredential | ProviderModelCredential | None = None
+
+            if credential_id:
+                if config_from == "predefined-model":
+                    credential_record = (
+                        db.session.query(ProviderCredential)
+                        .filter_by(
+                            id=credential_id,
+                            tenant_id=tenant_id,
+                            provider_name=provider_configuration.provider.provider,
+                        )
+                        .first()
+                    )
+                else:
+                    credential_record = (
+                        db.session.query(ProviderModelCredential)
+                        .filter_by(
+                            id=credential_id,
+                            tenant_id=tenant_id,
+                            provider_name=provider_configuration.provider.provider,
+                            model_name=model,
+                            model_type=model_type_enum.to_origin_model_type(),
+                        )
+                        .first()
+                    )
+                if not credential_record:
+                    raise ValueError(f"Provider credential with id {credential_id} not found")
+                name = credential_record.credential_name
+
+            if not name:
+                raise ValueError("Invalid load balancing config name")
+
+            if enabled is None:
+                raise ValueError("Invalid load balancing config enabled")
+
+            # is config exists
+            if config_id:
+                config_id = str(config_id)
+
+                if config_id not in current_load_balancing_configs_dict:
+                    raise ValueError(f"Invalid load balancing config id: {config_id}")
+
+                updated_config_ids.add(config_id)
+
+                load_balancing_config = current_load_balancing_configs_dict[config_id]
+
+                if credentials:
+                    if not isinstance(credentials, dict):
+                        raise ValueError("Invalid load balancing config credentials")
+
+                    # validate custom provider config
+                    credentials = self._custom_credentials_validate(
+                        tenant_id=tenant_id,
+                        provider_configuration=provider_configuration,
+                        model_type=model_type_enum,
+                        model=model,
+                        credentials=credentials,
+                        load_balancing_model_config=load_balancing_config,
+                        validate=False,
+                    )
+
+                    # update load balancing config
+                    load_balancing_config.encrypted_config = json.dumps(credentials)
+
+                load_balancing_config.name = name
+                load_balancing_config.enabled = enabled
+                load_balancing_config.updated_at = naive_utc_now()
+                db.session.commit()
+
+                self._clear_credentials_cache(tenant_id, config_id)
+            else:
+                # create load balancing config
+                if name == "__inherit__":
+                    raise ValueError("Invalid load balancing config name")
+
+                if credential_id:
+                    credential_source = "provider" if config_from == "predefined-model" else "custom_model"
+                    assert credential_record is not None
+                    load_balancing_model_config = LoadBalancingModelConfig(
+                        tenant_id=tenant_id,
+                        provider_name=provider_configuration.provider.provider,
+                        model_type=model_type_enum.to_origin_model_type(),
+                        model_name=model,
+                        name=credential_record.credential_name,
+                        encrypted_config=credential_record.encrypted_config,
+                        credential_id=credential_id,
+                        credential_source_type=credential_source,
+                    )
+                else:
+                    if not credentials:
+                        raise ValueError("Invalid load balancing config credentials")
+
+                    if not isinstance(credentials, dict):
+                        raise ValueError("Invalid load balancing config credentials")
+
+                    # validate custom provider config
+                    credentials = self._custom_credentials_validate(
+                        tenant_id=tenant_id,
+                        provider_configuration=provider_configuration,
+                        model_type=model_type_enum,
+                        model=model,
+                        credentials=credentials,
+                        validate=False,
+                    )
+
+                    # create load balancing config
+                    load_balancing_model_config = LoadBalancingModelConfig(
+                        tenant_id=tenant_id,
+                        provider_name=provider_configuration.provider.provider,
+                        model_type=model_type_enum.to_origin_model_type(),
+                        model_name=model,
+                        name=name,
+                        encrypted_config=json.dumps(credentials),
+                    )
+
+                db.session.add(load_balancing_model_config)
+                db.session.commit()
+
+        # get deleted config ids
+        deleted_config_ids = set(current_load_balancing_configs_dict.keys()) - updated_config_ids
+        for config_id in deleted_config_ids:
+            db.session.delete(current_load_balancing_configs_dict[config_id])
+            db.session.commit()
+
+            self._clear_credentials_cache(tenant_id, config_id)
+
+    def validate_load_balancing_credentials(
+        self,
+        tenant_id: str,
+        provider: str,
+        model: str,
+        model_type: str,
+        credentials: dict,
+        config_id: str | None = None,
+    ):
+        """
+        Validate load balancing credentials.
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model_type: model type
+        :param model: model name
+        :param credentials: credentials
+        :param config_id: load balancing config id
+        :return:
+        """
+        # Get all provider configurations of the current workspace
+        provider_configurations = self.provider_manager.get_configurations(tenant_id)
+
+        # Get provider configuration
+        provider_configuration = provider_configurations.get(provider)
+        if not provider_configuration:
+            raise ValueError(f"Provider {provider} does not exist.")
+
+        # Convert model type to ModelType
+        model_type_enum = ModelType.value_of(model_type)
+
+        load_balancing_model_config = None
+        if config_id:
+            # Get load balancing config
+            load_balancing_model_config = (
+                db.session.query(LoadBalancingModelConfig)
+                .where(
+                    LoadBalancingModelConfig.tenant_id == tenant_id,
+                    LoadBalancingModelConfig.provider_name == provider,
+                    LoadBalancingModelConfig.model_type == model_type_enum.to_origin_model_type(),
+                    LoadBalancingModelConfig.model_name == model,
+                    LoadBalancingModelConfig.id == config_id,
+                )
+                .first()
+            )
+
+            if not load_balancing_model_config:
+                raise ValueError(f"Load balancing config {config_id} does not exist.")
+
+        # Validate custom provider config
+        self._custom_credentials_validate(
+            tenant_id=tenant_id,
+            provider_configuration=provider_configuration,
+            model_type=model_type_enum,
+            model=model,
+            credentials=credentials,
+            load_balancing_model_config=load_balancing_model_config,
+        )
+
+    def _custom_credentials_validate(
+        self,
+        tenant_id: str,
+        provider_configuration: ProviderConfiguration,
+        model_type: ModelType,
+        model: str,
+        credentials: dict,
+        load_balancing_model_config: LoadBalancingModelConfig | None = None,
+        validate: bool = True,
+    ):
+        """
+        Validate custom credentials.
+        :param tenant_id: workspace id
+        :param provider_configuration: provider configuration
+        :param model_type: model type
+        :param model: model name
+        :param credentials: credentials
+        :param load_balancing_model_config: load balancing model config
+        :param validate: validate credentials
+        :return:
+        """
+        # Get credential form schemas from model credential schema or provider credential schema
+        credential_schemas = self._get_credential_schema(provider_configuration)
+
+        # Get provider credential secret variables
+        provider_credential_secret_variables = provider_configuration.extract_secret_variables(
+            credential_schemas.credential_form_schemas
+        )
+
+        if load_balancing_model_config:
+            try:
+                # fix origin data
+                if load_balancing_model_config.encrypted_config:
+                    original_credentials = json.loads(load_balancing_model_config.encrypted_config)
+                else:
+                    original_credentials = {}
+            except JSONDecodeError:
+                original_credentials = {}
+
+            # encrypt credentials
+            for key, value in credentials.items():
+                if key in provider_credential_secret_variables:
+                    # if send [__HIDDEN__] in secret input, it will be same as original value
+                    if value == HIDDEN_VALUE and key in original_credentials:
+                        credentials[key] = encrypter.decrypt_token(tenant_id, original_credentials[key])
+
+        if validate:
+            model_provider_factory = ModelProviderFactory(tenant_id)
+            if isinstance(credential_schemas, ModelCredentialSchema):
+                credentials = model_provider_factory.model_credentials_validate(
+                    provider=provider_configuration.provider.provider,
+                    model_type=model_type,
+                    model=model,
+                    credentials=credentials,
+                )
+            else:
+                credentials = model_provider_factory.provider_credentials_validate(
+                    provider=provider_configuration.provider.provider, credentials=credentials
+                )
+
+        for key, value in credentials.items():
+            if key in provider_credential_secret_variables:
+                credentials[key] = encrypter.encrypt_token(tenant_id, value)
+
+        return credentials
+
+    def _get_credential_schema(
+        self, provider_configuration: ProviderConfiguration
+    ) -> Union[ModelCredentialSchema, ProviderCredentialSchema]:
+        """Get form schemas."""
+        if provider_configuration.provider.model_credential_schema:
+            return provider_configuration.provider.model_credential_schema
+        elif provider_configuration.provider.provider_credential_schema:
+            return provider_configuration.provider.provider_credential_schema
+        else:
+            raise ValueError("No credential schema found")
+
+    def _clear_credentials_cache(self, tenant_id: str, config_id: str):
+        """
+        Clear credentials cache.
+        :param tenant_id: workspace id
+        :param config_id: load balancing config id
+        :return:
+        """
+        provider_model_credentials_cache = ProviderCredentialsCache(
+            tenant_id=tenant_id, identity_id=config_id, cache_type=ProviderCredentialsCacheType.LOAD_BALANCING_MODEL
+        )
+
+        provider_model_credentials_cache.delete()