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>
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user