tjdify/api/services/model_load_balancing_service.py

import datetime
import json
import logging
from json import JSONDecodeError
from typing import Optional

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 import model_provider_factory
from core.provider_manager import ProviderManager
from extensions.ext_database import db
from models.provider import LoadBalancingModelConfig

logger = logging.getLogger(__name__)


class ModelLoadBalancingService:

    def __init__(self) -> None:
        self.provider_manager = ProviderManager()

    def enable_model_load_balancing(self, tenant_id: str, provider: str, model: str, model_type: str) -> None:
        """
        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) -> None:
        """
        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) \
            -> 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 = ModelType.value_of(model_type)

        # Get provider model setting
        provider_model_setting = provider_configuration.get_provider_model_setting(
            model_type=model_type,
            model=model,
        )

        is_load_balancing_enabled = False
        if provider_model_setting and provider_model_setting.load_balancing_enabled:
            is_load_balancing_enabled = True

        # Get load balancing configurations
        load_balancing_configs = db.session.query(LoadBalancingModelConfig) \
            .filter(
            LoadBalancingModelConfig.tenant_id == tenant_id,
            LoadBalancingModelConfig.provider_name == provider_configuration.provider.provider,
            LoadBalancingModelConfig.model_type == model_type.to_origin_model_type(),
            LoadBalancingModelConfig.model_name == model
        ).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)

                # 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,
                config_id=load_balancing_config.id
            )

            try:
                if load_balancing_config.encrypted_config:
                    credentials = 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:
                        credentials[variable] = encrypter.decrypt_token_with_decoding(
                            credentials.get(variable),
                            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,
                '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) \
            -> Optional[dict]:
        """
        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 = ModelType.value_of(model_type)

        # Get load balancing configurations
        load_balancing_model_config = db.session.query(LoadBalancingModelConfig) \
            .filter(
            LoadBalancingModelConfig.tenant_id == tenant_id,
            LoadBalancingModelConfig.provider_name == provider_configuration.provider.provider,
            LoadBalancingModelConfig.model_type == model_type.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]) -> None:
        """
        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
        :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 = ModelType.value_of(model_type)

        if not isinstance(configs, list):
            raise ValueError('Invalid load balancing configs')

        current_load_balancing_configs = db.session.query(LoadBalancingModelConfig) \
            .filter(
            LoadBalancingModelConfig.tenant_id == tenant_id,
            LoadBalancingModelConfig.provider_name == provider_configuration.provider.provider,
            LoadBalancingModelConfig.model_type == model_type.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')
            enabled = config.get('enabled')

            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('Invalid load balancing config id: {}'.format(config_id))

                updated_config_ids.add(config_id)

                load_balancing_config = current_load_balancing_configs_dict[config_id]

                # check duplicate name
                for current_load_balancing_config in current_load_balancing_configs:
                    if current_load_balancing_config.id != config_id and current_load_balancing_config.name == name:
                        raise ValueError('Load balancing config name {} already exists'.format(name))

                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,
                        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 = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
                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')

                # check duplicate name
                for current_load_balancing_config in current_load_balancing_configs:
                    if current_load_balancing_config.name == name:
                        raise ValueError('Load balancing config name {} already exists'.format(name))

                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,
                    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.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: Optional[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 = 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) \
                .filter(
                LoadBalancingModelConfig.tenant_id == tenant_id,
                LoadBalancingModelConfig.provider_name == provider,
                LoadBalancingModelConfig.model_type == model_type.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,
            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: Optional[LoadBalancingModelConfig] = None,
                                     validate: bool = True) -> dict:
        """
        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__]' and key in original_credentials:
                        credentials[key] = encrypter.decrypt_token(tenant_id, original_credentials[key])

        if validate:
            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) \
            -> ModelCredentialSchema | ProviderCredentialSchema:
        """
        Get form schemas.
        :param provider_configuration: provider configuration
        :return:
        """
        # Get credential form schemas from model credential schema or provider credential schema
        if provider_configuration.provider.model_credential_schema:
            credential_schema = provider_configuration.provider.model_credential_schema
        else:
            credential_schema = provider_configuration.provider.provider_credential_schema

        return credential_schema

    def _clear_credentials_cache(self, tenant_id: str, config_id: str) -> None:
        """
        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()