feat: backend model load balancing support (#4927)

api/services/dataset_service.py

@@ -4,7 +4,7 @@ import logging
 import random
 import time
 import uuid
-from typing import Optional, cast
+from typing import Optional
 
 from flask import current_app
 from flask_login import current_user
@@ -13,7 +13,6 @@ from sqlalchemy import func
 from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.model_manager import ModelManager
 from core.model_runtime.entities.model_entities import ModelType
-from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
 from core.rag.datasource.keyword.keyword_factory import Keyword
 from core.rag.models.document import Document as RAGDocument
 from events.dataset_event import dataset_was_deleted
@@ -1144,10 +1143,7 @@ class SegmentService:
                 model=dataset.embedding_model
             )
             # calc embedding use tokens
-            model_type_instance = cast(TextEmbeddingModel, embedding_model.model_type_instance)
-            tokens = model_type_instance.get_num_tokens(
-                model=embedding_model.model,
-                credentials=embedding_model.credentials,
+            tokens = embedding_model.get_text_embedding_num_tokens(
                 texts=[content]
             )
         lock_name = 'add_segment_lock_document_id_{}'.format(document.id)
@@ -1215,10 +1211,7 @@ class SegmentService:
                 tokens = 0
                 if dataset.indexing_technique == 'high_quality' and embedding_model:
                     # calc embedding use tokens
-                    model_type_instance = cast(TextEmbeddingModel, embedding_model.model_type_instance)
-                    tokens = model_type_instance.get_num_tokens(
-                        model=embedding_model.model,
-                        credentials=embedding_model.credentials,
+                    tokens = embedding_model.get_text_embedding_num_tokens(
                         texts=[content]
                     )
                 segment_document = DocumentSegment(
@@ -1321,10 +1314,7 @@ class SegmentService:
                     )
 
                     # calc embedding use tokens
-                    model_type_instance = cast(TextEmbeddingModel, embedding_model.model_type_instance)
-                    tokens = model_type_instance.get_num_tokens(
-                        model=embedding_model.model,
-                        credentials=embedding_model.credentials,
+                    tokens = embedding_model.get_text_embedding_num_tokens(
                         texts=[content]
                     )
                 segment.content = content

api/services/entities/model_provider_entities.py

@@ -4,10 +4,10 @@ from typing import Optional
 from flask import current_app
 from pydantic import BaseModel
 
-from core.entities.model_entities import ModelStatus, ModelWithProviderEntity
+from core.entities.model_entities import ModelWithProviderEntity, ProviderModelWithStatusEntity
 from core.entities.provider_entities import QuotaConfiguration
 from core.model_runtime.entities.common_entities import I18nObject
-from core.model_runtime.entities.model_entities import ModelType, ProviderModel
+from core.model_runtime.entities.model_entities import ModelType
 from core.model_runtime.entities.provider_entities import (
     ConfigurateMethod,
     ModelCredentialSchema,
@@ -79,13 +79,6 @@ class ProviderResponse(BaseModel):
             )
 
 
-class ModelResponse(ProviderModel):
-    """
-    Model class for model response.
-    """
-    status: ModelStatus
-
-
 class ProviderWithModelsResponse(BaseModel):
     """
     Model class for provider with models response.
@@ -95,7 +88,7 @@ class ProviderWithModelsResponse(BaseModel):
     icon_small: Optional[I18nObject] = None
     icon_large: Optional[I18nObject] = None
     status: CustomConfigurationStatus
-    models: list[ModelResponse]
+    models: list[ProviderModelWithStatusEntity]
 
     def __init__(self, **data) -> None:
         super().__init__(**data)

api/services/feature_service.py

@@ -29,6 +29,7 @@ class FeatureModel(BaseModel):
     documents_upload_quota: LimitationModel = LimitationModel(size=0, limit=50)
     docs_processing: str = 'standard'
     can_replace_logo: bool = False
+    model_load_balancing_enabled: bool = False
 
 
 class SystemFeatureModel(BaseModel):
@@ -63,6 +64,7 @@ class FeatureService:
     @classmethod
     def _fulfill_params_from_env(cls, features: FeatureModel):
         features.can_replace_logo = current_app.config['CAN_REPLACE_LOGO']
+        features.model_load_balancing_enabled = current_app.config['MODEL_LB_ENABLED']
 
     @classmethod
     def _fulfill_params_from_billing_api(cls, features: FeatureModel, tenant_id: str):
@@ -72,23 +74,34 @@ class FeatureService:
         features.billing.subscription.plan = billing_info['subscription']['plan']
         features.billing.subscription.interval = billing_info['subscription']['interval']
 
-        features.members.size = billing_info['members']['size']
-        features.members.limit = billing_info['members']['limit']
+        if 'members' in billing_info:
+            features.members.size = billing_info['members']['size']
+            features.members.limit = billing_info['members']['limit']
 
-        features.apps.size = billing_info['apps']['size']
-        features.apps.limit = billing_info['apps']['limit']
+        if 'apps' in billing_info:
+            features.apps.size = billing_info['apps']['size']
+            features.apps.limit = billing_info['apps']['limit']
 
-        features.vector_space.size = billing_info['vector_space']['size']
-        features.vector_space.limit = billing_info['vector_space']['limit']
+        if 'vector_space' in billing_info:
+            features.vector_space.size = billing_info['vector_space']['size']
+            features.vector_space.limit = billing_info['vector_space']['limit']
 
-        features.documents_upload_quota.size = billing_info['documents_upload_quota']['size']
-        features.documents_upload_quota.limit = billing_info['documents_upload_quota']['limit']
+        if 'documents_upload_quota' in billing_info:
+            features.documents_upload_quota.size = billing_info['documents_upload_quota']['size']
+            features.documents_upload_quota.limit = billing_info['documents_upload_quota']['limit']
 
-        features.annotation_quota_limit.size = billing_info['annotation_quota_limit']['size']
-        features.annotation_quota_limit.limit = billing_info['annotation_quota_limit']['limit']
+        if 'annotation_quota_limit' in billing_info:
+            features.annotation_quota_limit.size = billing_info['annotation_quota_limit']['size']
+            features.annotation_quota_limit.limit = billing_info['annotation_quota_limit']['limit']
 
-        features.docs_processing = billing_info['docs_processing']
-        features.can_replace_logo = billing_info['can_replace_logo']
+        if 'docs_processing' in billing_info:
+            features.docs_processing = billing_info['docs_processing']
+
+        if 'can_replace_logo' in billing_info:
+            features.can_replace_logo = billing_info['can_replace_logo']
+
+        if 'model_load_balancing_enabled' in billing_info:
+            features.model_load_balancing_enabled = billing_info['model_load_balancing_enabled']
 
     @classmethod
     def _fulfill_params_from_enterprise(cls, features):

api/services/model_load_balancing_service.py

@@ -0,0 +1,565 @@
+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()

api/services/model_provider_service.py

@@ -6,7 +6,7 @@ from typing import Optional, cast
 import requests
 from flask import current_app
 
-from core.entities.model_entities import ModelStatus
+from core.entities.model_entities import ModelStatus, ProviderModelWithStatusEntity
 from core.model_runtime.entities.model_entities import ModelType, ParameterRule
 from core.model_runtime.model_providers import model_provider_factory
 from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
@@ -16,7 +16,6 @@ from services.entities.model_provider_entities import (
     CustomConfigurationResponse,
     CustomConfigurationStatus,
     DefaultModelResponse,
-    ModelResponse,
     ModelWithProviderEntityResponse,
     ProviderResponse,
     ProviderWithModelsResponse,
@@ -303,6 +302,9 @@ class ModelProviderService:
             if model.deprecated:
                 continue
 
+            if model.status != ModelStatus.ACTIVE:
+                continue
+
             provider_models[model.provider.provider].append(model)
 
         # convert to ProviderWithModelsResponse list
@@ -313,24 +315,22 @@ class ModelProviderService:
 
             first_model = models[0]
 
-            has_active_models = any([model.status == ModelStatus.ACTIVE for model in models])
-
             providers_with_models.append(
                 ProviderWithModelsResponse(
                     provider=provider,
                     label=first_model.provider.label,
                     icon_small=first_model.provider.icon_small,
                     icon_large=first_model.provider.icon_large,
-                    status=CustomConfigurationStatus.ACTIVE
-                    if has_active_models else CustomConfigurationStatus.NO_CONFIGURE,
-                    models=[ModelResponse(
+                    status=CustomConfigurationStatus.ACTIVE,
+                    models=[ProviderModelWithStatusEntity(
                         model=model.model,
                         label=model.label,
                         model_type=model.model_type,
                         features=model.features,
                         fetch_from=model.fetch_from,
                         model_properties=model.model_properties,
-                        status=model.status
+                        status=model.status,
+                        load_balancing_enabled=model.load_balancing_enabled
                     ) for model in models]
                 )
             )
@@ -486,6 +486,54 @@ class ModelProviderService:
         # Switch preferred provider type
         provider_configuration.switch_preferred_provider_type(preferred_provider_type_enum)
 
+    def enable_model(self, tenant_id: str, provider: str, model: str, model_type: str) -> None:
+        """
+        enable model.
+
+        :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
+        provider_configuration.enable_model(
+            model=model,
+            model_type=ModelType.value_of(model_type)
+        )
+
+    def disable_model(self, tenant_id: str, provider: str, model: str, model_type: str) -> None:
+        """
+        disable model.
+
+        :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
+        provider_configuration.disable_model(
+            model=model,
+            model_type=ModelType.value_of(model_type)
+        )
+
     def free_quota_submit(self, tenant_id: str, provider: str):
         api_key = os.environ.get("FREE_QUOTA_APPLY_API_KEY")
         api_base_url = os.environ.get("FREE_QUOTA_APPLY_BASE_URL")

api/services/workflow_service.py

@@ -68,7 +68,7 @@ class WorkflowService:
                             account: Account) -> Workflow:
         """
         Sync draft workflow
-        @throws WorkflowHashNotEqualError
+        :raises WorkflowHashNotEqualError
         """
         # fetch draft workflow by app_model
         workflow = self.get_draft_workflow(app_model=app_model)