feat: add multi model credentials (#24451)

Co-authored-by: zxhlyh <jasonapring2015@outlook.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/services/entities/model_provider_entities.py

@@ -8,7 +8,12 @@ from core.entities.model_entities import (
     ModelWithProviderEntity,
     ProviderModelWithStatusEntity,
 )
-from core.entities.provider_entities import ProviderQuotaType, QuotaConfiguration
+from core.entities.provider_entities import (
+    CredentialConfiguration,
+    CustomModelConfiguration,
+    ProviderQuotaType,
+    QuotaConfiguration,
+)
 from core.model_runtime.entities.common_entities import I18nObject
 from core.model_runtime.entities.model_entities import ModelType
 from core.model_runtime.entities.provider_entities import (
@@ -36,6 +41,10 @@ class CustomConfigurationResponse(BaseModel):
     """
 
     status: CustomConfigurationStatus
+    current_credential_id: Optional[str] = None
+    current_credential_name: Optional[str] = None
+    available_credentials: Optional[list[CredentialConfiguration]] = None
+    custom_models: Optional[list[CustomModelConfiguration]] = None
 
 
 class SystemConfigurationResponse(BaseModel):

api/services/errors/app_model_config.py

@@ -3,3 +3,7 @@ from services.errors.base import BaseServiceError
 
 class AppModelConfigBrokenError(BaseServiceError):
     pass
+
+
+class ProviderNotFoundError(BaseServiceError):
+    pass

api/services/model_load_balancing_service.py

@@ -17,7 +17,7 @@ from core.model_runtime.model_providers.model_provider_factory import ModelProvi
 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
+from models.provider import LoadBalancingModelConfig, ProviderCredential, ProviderModelCredential
 
 logger = logging.getLogger(__name__)
 
@@ -185,6 +185,7 @@ class ModelLoadBalancingService:
                     "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,
@@ -280,7 +281,7 @@ class ModelLoadBalancingService:
         return inherit_config
 
     def update_load_balancing_configs(
-        self, tenant_id: str, provider: str, model: str, model_type: str, configs: list[dict]
+        self, tenant_id: str, provider: str, model: str, model_type: str, configs: list[dict], config_from: str
     ) -> None:
         """
         Update load balancing configurations.
@@ -289,6 +290,7 @@ class ModelLoadBalancingService:
         :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
@@ -327,8 +329,37 @@ class ModelLoadBalancingService:
             config_id = config.get("id")
             name = config.get("name")
             credentials = config.get("credentials")
+            credential_id = config.get("credential_id")
             enabled = config.get("enabled")
 
+            if credential_id:
+                credential_record: ProviderCredential | ProviderModelCredential | None = None
+                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")
 
@@ -346,11 +377,6 @@ class ModelLoadBalancingService:
 
                 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(f"Load balancing config name {name} already exists")
-
                 if credentials:
                     if not isinstance(credentials, dict):
                         raise ValueError("Invalid load balancing config credentials")
@@ -377,39 +403,48 @@ class ModelLoadBalancingService:
                 self._clear_credentials_cache(tenant_id, config_id)
             else:
                 # create load balancing config
-                if name == "__inherit__":
+                if name in {"__inherit__", "__delete__"}:
                     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(f"Load balancing config name {name} already exists")
+                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 credentials:
-                    raise ValueError("Invalid load balancing config credentials")
+                    if not isinstance(credentials, dict):
+                        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,
+                    )
 
-                # 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),
-                )
+                    # 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()

api/services/model_provider_service.py

@@ -16,6 +16,7 @@ from services.entities.model_provider_entities import (
     SimpleProviderEntityResponse,
     SystemConfigurationResponse,
 )
+from services.errors.app_model_config import ProviderNotFoundError
 
 logger = logging.getLogger(__name__)
 
@@ -28,6 +29,29 @@ class ModelProviderService:
     def __init__(self) -> None:
         self.provider_manager = ProviderManager()
 
+    def _get_provider_configuration(self, tenant_id: str, provider: str):
+        """
+        Get provider configuration or raise exception if not found.
+
+        Args:
+            tenant_id: Workspace identifier
+            provider: Provider name
+
+        Returns:
+            Provider configuration instance
+
+        Raises:
+            ProviderNotFoundError: If provider doesn't exist
+        """
+        # Get all provider configurations of the current workspace
+        provider_configurations = self.provider_manager.get_configurations(tenant_id)
+        provider_configuration = provider_configurations.get(provider)
+
+        if not provider_configuration:
+            raise ProviderNotFoundError(f"Provider {provider} does not exist.")
+
+        return provider_configuration
+
     def get_provider_list(self, tenant_id: str, model_type: Optional[str] = None) -> list[ProviderResponse]:
         """
         get provider list.
@@ -46,6 +70,9 @@ class ModelProviderService:
                 if model_type_entity not in provider_configuration.provider.supported_model_types:
                     continue
 
+            provider_config = provider_configuration.custom_configuration.provider
+            model_config = provider_configuration.custom_configuration.models
+
             provider_response = ProviderResponse(
                 tenant_id=tenant_id,
                 provider=provider_configuration.provider.provider,
@@ -63,7 +90,11 @@ class ModelProviderService:
                 custom_configuration=CustomConfigurationResponse(
                     status=CustomConfigurationStatus.ACTIVE
                     if provider_configuration.is_custom_configuration_available()
-                    else CustomConfigurationStatus.NO_CONFIGURE
+                    else CustomConfigurationStatus.NO_CONFIGURE,
+                    current_credential_id=getattr(provider_config, "current_credential_id", None),
+                    current_credential_name=getattr(provider_config, "current_credential_name", None),
+                    available_credentials=getattr(provider_config, "available_credentials", []),
+                    custom_models=model_config,
                 ),
                 system_configuration=SystemConfigurationResponse(
                     enabled=provider_configuration.system_configuration.enabled,
@@ -82,8 +113,8 @@ class ModelProviderService:
         For the model provider page,
         only supports passing in a single provider to query the list of supported models.
 
-        :param tenant_id:
-        :param provider:
+        :param tenant_id: workspace id
+        :param provider: provider name
         :return:
         """
         # Get all provider configurations of the current workspace
@@ -95,98 +126,111 @@ class ModelProviderService:
             for model in provider_configurations.get_models(provider=provider)
         ]
 
-    def get_provider_credentials(self, tenant_id: str, provider: str) -> Optional[dict]:
+    def get_provider_credential(
+        self, tenant_id: str, provider: str, credential_id: Optional[str] = None
+    ) -> Optional[dict]:
         """
         get provider credentials.
-        """
-        provider_configurations = self.provider_manager.get_configurations(tenant_id)
-        provider_configuration = provider_configurations.get(provider)
-        if not provider_configuration:
-            raise ValueError(f"Provider {provider} does not exist.")
-
-        return provider_configuration.get_custom_credentials(obfuscated=True)
-
-    def provider_credentials_validate(self, tenant_id: str, provider: str, credentials: dict) -> None:
-        """
-        validate provider credentials.
-
-        :param tenant_id:
-        :param provider:
-        :param credentials:
-        """
-        # 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.")
-
-        provider_configuration.custom_credentials_validate(credentials)
-
-    def save_provider_credentials(self, tenant_id: str, provider: str, credentials: dict) -> None:
-        """
-        save custom provider config.
 
         :param tenant_id: workspace id
         :param provider: provider name
-        :param credentials: provider credentials
+        :param credential_id: credential id, if not provided, return current used credentials
         :return:
         """
-        # Get all provider configurations of the current workspace
-        provider_configurations = self.provider_manager.get_configurations(tenant_id)
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        return provider_configuration.get_provider_credential(credential_id=credential_id)  # type: ignore
 
-        # Get provider configuration
-        provider_configuration = provider_configurations.get(provider)
-        if not provider_configuration:
-            raise ValueError(f"Provider {provider} does not exist.")
-
-        # Add or update custom provider credentials.
-        provider_configuration.add_or_update_custom_credentials(credentials)
-
-    def remove_provider_credentials(self, tenant_id: str, provider: str) -> None:
+    def validate_provider_credentials(self, tenant_id: str, provider: str, credentials: dict) -> None:
         """
-        remove custom provider config.
+        validate provider credentials before saving.
 
         :param tenant_id: workspace id
         :param provider: provider name
+        :param credentials: provider credentials dict
+        """
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.validate_provider_credentials(credentials)
+
+    def create_provider_credential(
+        self, tenant_id: str, provider: str, credentials: dict, credential_name: str
+    ) -> None:
+        """
+        Create and save new provider credentials.
+
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param credentials: provider credentials dict
+        :param credential_name: credential name
         :return:
         """
-        # Get all provider configurations of the current workspace
-        provider_configurations = self.provider_manager.get_configurations(tenant_id)
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.create_provider_credential(credentials, credential_name)
 
-        # Get provider configuration
-        provider_configuration = provider_configurations.get(provider)
-        if not provider_configuration:
-            raise ValueError(f"Provider {provider} does not exist.")
-
-        # Remove custom provider credentials.
-        provider_configuration.delete_custom_credentials()
-
-    def get_model_credentials(self, tenant_id: str, provider: str, model_type: str, model: str) -> Optional[dict]:
+    def update_provider_credential(
+        self,
+        tenant_id: str,
+        provider: str,
+        credentials: dict,
+        credential_id: str,
+        credential_name: str,
+    ) -> None:
         """
-        get model credentials.
+        update a saved provider credential (by credential_id).
+
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param credentials: provider credentials dict
+        :param credential_id: credential id
+        :param credential_name: credential name
+        :return:
+        """
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.update_provider_credential(
+            credential_id=credential_id,
+            credentials=credentials,
+            credential_name=credential_name,
+        )
+
+    def remove_provider_credential(self, tenant_id: str, provider: str, credential_id: str) -> None:
+        """
+        remove a saved provider credential (by credential_id).
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param credential_id: credential id
+        :return:
+        """
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.delete_provider_credential(credential_id=credential_id)
+
+    def switch_active_provider_credential(self, tenant_id: str, provider: str, credential_id: str) -> None:
+        """
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param credential_id: credential id
+        :return:
+        """
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.switch_active_provider_credential(credential_id=credential_id)
+
+    def get_model_credential(
+        self, tenant_id: str, provider: str, model_type: str, model: str, credential_id: str | None
+    ) -> Optional[dict]:
+        """
+        Retrieve model-specific credentials.
 
         :param tenant_id: workspace id
         :param provider: provider name
         :param model_type: model type
         :param model: model name
+        :param credential_id: Optional credential ID, uses current if not provided
         :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.")
-
-        # Get model custom credentials from ProviderModel if exists
-        return provider_configuration.get_custom_model_credentials(
-            model_type=ModelType.value_of(model_type), model=model, obfuscated=True
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        return provider_configuration.get_custom_model_credential(  # type: ignore
+            model_type=ModelType.value_of(model_type), model=model, credential_id=credential_id
         )
 
-    def model_credentials_validate(
+    def validate_model_credentials(
         self, tenant_id: str, provider: str, model_type: str, model: str, credentials: dict
     ) -> None:
         """
@@ -196,49 +240,122 @@ class ModelProviderService:
         :param provider: provider name
         :param model_type: model type
         :param model: model name
-        :param credentials: model credentials
+        :param credentials: model credentials dict
         :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.")
-
-        # Validate model credentials
-        provider_configuration.custom_model_credentials_validate(
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.validate_custom_model_credentials(
             model_type=ModelType.value_of(model_type), model=model, credentials=credentials
         )
 
-    def save_model_credentials(
-        self, tenant_id: str, provider: str, model_type: str, model: str, credentials: dict
+    def create_model_credential(
+        self, tenant_id: str, provider: str, model_type: str, model: str, credentials: dict, credential_name: str
     ) -> None:
         """
-        save model credentials.
+        create and save model credentials.
 
         :param tenant_id: workspace id
         :param provider: provider name
         :param model_type: model type
         :param model: model name
-        :param credentials: model credentials
+        :param credentials: model credentials dict
+        :param credential_name: credential name
         :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.")
-
-        # Add or update custom model credentials
-        provider_configuration.add_or_update_custom_model_credentials(
-            model_type=ModelType.value_of(model_type), model=model, credentials=credentials
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.create_custom_model_credential(
+            model_type=ModelType.value_of(model_type),
+            model=model,
+            credentials=credentials,
+            credential_name=credential_name,
         )
 
-    def remove_model_credentials(self, tenant_id: str, provider: str, model_type: str, model: str) -> None:
+    def update_model_credential(
+        self,
+        tenant_id: str,
+        provider: str,
+        model_type: str,
+        model: str,
+        credentials: dict,
+        credential_id: str,
+        credential_name: str,
+    ) -> None:
+        """
+        update model credentials.
+
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model_type: model type
+        :param model: model name
+        :param credentials: model credentials dict
+        :param credential_id: credential id
+        :param credential_name: credential name
+        :return:
+        """
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.update_custom_model_credential(
+            model_type=ModelType.value_of(model_type),
+            model=model,
+            credentials=credentials,
+            credential_id=credential_id,
+            credential_name=credential_name,
+        )
+
+    def remove_model_credential(
+        self, tenant_id: str, provider: str, model_type: str, model: str, credential_id: str
+    ) -> None:
+        """
+        remove model credentials.
+
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model_type: model type
+        :param model: model name
+        :param credential_id: credential id
+        :return:
+        """
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.delete_custom_model_credential(
+            model_type=ModelType.value_of(model_type), model=model, credential_id=credential_id
+        )
+
+    def switch_active_custom_model_credential(
+        self, tenant_id: str, provider: str, model_type: str, model: str, credential_id: str
+    ) -> None:
+        """
+        switch model credentials.
+
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model_type: model type
+        :param model: model name
+        :param credential_id: credential id
+        :return:
+        """
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.switch_custom_model_credential(
+            model_type=ModelType.value_of(model_type), model=model, credential_id=credential_id
+        )
+
+    def add_model_credential_to_model_list(
+        self, tenant_id: str, provider: str, model_type: str, model: str, credential_id: str
+    ) -> None:
+        """
+        add model credentials to model list.
+
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model_type: model type
+        :param model: model name
+        :param credential_id: credential id
+        :return:
+        """
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.add_model_credential_to_model(
+            model_type=ModelType.value_of(model_type), model=model, credential_id=credential_id
+        )
+
+    def remove_model(self, tenant_id: str, provider: str, model_type: str, model: str) -> None:
         """
         remove model credentials.
 
@@ -248,16 +365,8 @@ class ModelProviderService:
         :param model: model name
         :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.")
-
-        # Remove custom model credentials
-        provider_configuration.delete_custom_model_credentials(model_type=ModelType.value_of(model_type), model=model)
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
+        provider_configuration.delete_custom_model(model_type=ModelType.value_of(model_type), model=model)
 
     def get_models_by_model_type(self, tenant_id: str, model_type: str) -> list[ProviderWithModelsResponse]:
         """
@@ -331,13 +440,7 @@ class ModelProviderService:
         :param model: model name
         :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.")
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
 
         # fetch credentials
         credentials = provider_configuration.get_current_credentials(model_type=ModelType.LLM, model=model)
@@ -424,17 +527,11 @@ class ModelProviderService:
         :param preferred_provider_type: preferred provider type
         :return:
         """
-        # Get all provider configurations of the current workspace
-        provider_configurations = self.provider_manager.get_configurations(tenant_id)
+        provider_configuration = self._get_provider_configuration(tenant_id, provider)
 
         # Convert preferred_provider_type to ProviderType
         preferred_provider_type_enum = ProviderType.value_of(preferred_provider_type)
 
-        # Get provider configuration
-        provider_configuration = provider_configurations.get(provider)
-        if not provider_configuration:
-            raise ValueError(f"Provider {provider} does not exist.")
-
         # Switch preferred provider type
         provider_configuration.switch_preferred_provider_type(preferred_provider_type_enum)
 
@@ -448,15 +545,7 @@ class ModelProviderService:
         :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 = self._get_provider_configuration(tenant_id, provider)
         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:
@@ -469,13 +558,5 @@ class ModelProviderService:
         :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 = self._get_provider_configuration(tenant_id, provider)
         provider_configuration.disable_model(model=model, model_type=ModelType.value_of(model_type))