feat: backend model load balancing support (#4927)
takatost
2 years ago
committed by
GitHub
parent
52ec152dd3
commit
d1dbbc1e33
@ -0,0 +1,106 @@
|
||||
from flask_restful import Resource, reqparse
|
||||
from werkzeug.exceptions import Forbidden
|
||||
|
||||
from controllers.console import api
|
||||
from controllers.console.setup import setup_required
|
||||
from controllers.console.wraps import account_initialization_required
|
||||
from core.model_runtime.entities.model_entities import ModelType
|
||||
from core.model_runtime.errors.validate import CredentialsValidateFailedError
|
||||
from libs.login import current_user, login_required
|
||||
from models.account import TenantAccountRole
|
||||
from services.model_load_balancing_service import ModelLoadBalancingService
|
||||
|
||||
|
||||
class LoadBalancingCredentialsValidateApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
def post(self, provider: str):
|
||||
if not TenantAccountRole.is_privileged_role(current_user.current_tenant.current_role):
|
||||
raise Forbidden()
|
||||
|
||||
tenant_id = current_user.current_tenant_id
|
||||
|
||||
parser = reqparse.RequestParser()
|
||||
parser.add_argument('model', type=str, required=True, nullable=False, location='json')
|
||||
parser.add_argument('model_type', type=str, required=True, nullable=False,
|
||||
choices=[mt.value for mt in ModelType], location='json')
|
||||
parser.add_argument('credentials', type=dict, required=True, nullable=False, location='json')
|
||||
args = parser.parse_args()
|
||||
|
||||
# validate model load balancing credentials
|
||||
model_load_balancing_service = ModelLoadBalancingService()
|
||||
|
||||
result = True
|
||||
error = None
|
||||
|
||||
try:
|
||||
model_load_balancing_service.validate_load_balancing_credentials(
|
||||
tenant_id=tenant_id,
|
||||
provider=provider,
|
||||
model=args['model'],
|
||||
model_type=args['model_type'],
|
||||
credentials=args['credentials']
|
||||
)
|
||||
except CredentialsValidateFailedError as ex:
|
||||
result = False
|
||||
error = str(ex)
|
||||
|
||||
response = {'result': 'success' if result else 'error'}
|
||||
|
||||
if not result:
|
||||
response['error'] = error
|
||||
|
||||
return response
|
||||
|
||||
|
||||
class LoadBalancingConfigCredentialsValidateApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
def post(self, provider: str, config_id: str):
|
||||
if not TenantAccountRole.is_privileged_role(current_user.current_tenant.current_role):
|
||||
raise Forbidden()
|
||||
|
||||
tenant_id = current_user.current_tenant_id
|
||||
|
||||
parser = reqparse.RequestParser()
|
||||
parser.add_argument('model', type=str, required=True, nullable=False, location='json')
|
||||
parser.add_argument('model_type', type=str, required=True, nullable=False,
|
||||
choices=[mt.value for mt in ModelType], location='json')
|
||||
parser.add_argument('credentials', type=dict, required=True, nullable=False, location='json')
|
||||
args = parser.parse_args()
|
||||
|
||||
# validate model load balancing config credentials
|
||||
model_load_balancing_service = ModelLoadBalancingService()
|
||||
|
||||
result = True
|
||||
error = None
|
||||
|
||||
try:
|
||||
model_load_balancing_service.validate_load_balancing_credentials(
|
||||
tenant_id=tenant_id,
|
||||
provider=provider,
|
||||
model=args['model'],
|
||||
model_type=args['model_type'],
|
||||
credentials=args['credentials'],
|
||||
config_id=config_id,
|
||||
)
|
||||
except CredentialsValidateFailedError as ex:
|
||||
result = False
|
||||
error = str(ex)
|
||||
|
||||
response = {'result': 'success' if result else 'error'}
|
||||
|
||||
if not result:
|
||||
response['error'] = error
|
||||
|
||||
return response
|
||||
|
||||
|
||||
# Load Balancing Config
|
||||
api.add_resource(LoadBalancingCredentialsValidateApi,
|
||||
'/workspaces/current/model-providers/<string:provider>/models/load-balancing-configs/credentials-validate')
|
||||
|
||||
api.add_resource(LoadBalancingConfigCredentialsValidateApi,
|
||||
'/workspaces/current/model-providers/<string:provider>/models/load-balancing-configs/<string:config_id>/credentials-validate')
|
||||
@ -0,0 +1,126 @@
|
||||
"""add load balancing
|
||||
|
||||
Revision ID: 4e99a8df00ff
|
||||
Revises: 47cc7df8c4f3
|
||||
Create Date: 2024-05-10 12:08:09.812736
|
||||
|
||||
"""
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
import models as models
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = '4e99a8df00ff'
|
||||
down_revision = '64a70a7aab8b'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table('load_balancing_model_configs',
|
||||
sa.Column('id', models.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
|
||||
sa.Column('tenant_id', models.StringUUID(), nullable=False),
|
||||
sa.Column('provider_name', sa.String(length=255), nullable=False),
|
||||
sa.Column('model_name', sa.String(length=255), nullable=False),
|
||||
sa.Column('model_type', sa.String(length=40), nullable=False),
|
||||
sa.Column('name', sa.String(length=255), nullable=False),
|
||||
sa.Column('encrypted_config', sa.Text(), nullable=True),
|
||||
sa.Column('enabled', sa.Boolean(), server_default=sa.text('true'), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP(0)'), nullable=False),
|
||||
sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP(0)'), nullable=False),
|
||||
sa.PrimaryKeyConstraint('id', name='load_balancing_model_config_pkey')
|
||||
)
|
||||
with op.batch_alter_table('load_balancing_model_configs', schema=None) as batch_op:
|
||||
batch_op.create_index('load_balancing_model_config_tenant_provider_model_idx', ['tenant_id', 'provider_name', 'model_type'], unique=False)
|
||||
|
||||
op.create_table('provider_model_settings',
|
||||
sa.Column('id', models.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
|
||||
sa.Column('tenant_id', models.StringUUID(), nullable=False),
|
||||
sa.Column('provider_name', sa.String(length=255), nullable=False),
|
||||
sa.Column('model_name', sa.String(length=255), nullable=False),
|
||||
sa.Column('model_type', sa.String(length=40), nullable=False),
|
||||
sa.Column('enabled', sa.Boolean(), server_default=sa.text('true'), nullable=False),
|
||||
sa.Column('load_balancing_enabled', sa.Boolean(), server_default=sa.text('false'), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP(0)'), nullable=False),
|
||||
sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP(0)'), nullable=False),
|
||||
sa.PrimaryKeyConstraint('id', name='provider_model_setting_pkey')
|
||||
)
|
||||
with op.batch_alter_table('provider_model_settings', schema=None) as batch_op:
|
||||
batch_op.create_index('provider_model_setting_tenant_provider_model_idx', ['tenant_id', 'provider_name', 'model_type'], unique=False)
|
||||
|
||||
with op.batch_alter_table('provider_models', schema=None) as batch_op:
|
||||
batch_op.alter_column('provider_name',
|
||||
existing_type=sa.VARCHAR(length=40),
|
||||
type_=sa.String(length=255),
|
||||
existing_nullable=False)
|
||||
|
||||
with op.batch_alter_table('provider_orders', schema=None) as batch_op:
|
||||
batch_op.alter_column('provider_name',
|
||||
existing_type=sa.VARCHAR(length=40),
|
||||
type_=sa.String(length=255),
|
||||
existing_nullable=False)
|
||||
|
||||
with op.batch_alter_table('providers', schema=None) as batch_op:
|
||||
batch_op.alter_column('provider_name',
|
||||
existing_type=sa.VARCHAR(length=40),
|
||||
type_=sa.String(length=255),
|
||||
existing_nullable=False)
|
||||
|
||||
with op.batch_alter_table('tenant_default_models', schema=None) as batch_op:
|
||||
batch_op.alter_column('provider_name',
|
||||
existing_type=sa.VARCHAR(length=40),
|
||||
type_=sa.String(length=255),
|
||||
existing_nullable=False)
|
||||
|
||||
with op.batch_alter_table('tenant_preferred_model_providers', schema=None) as batch_op:
|
||||
batch_op.alter_column('provider_name',
|
||||
existing_type=sa.VARCHAR(length=40),
|
||||
type_=sa.String(length=255),
|
||||
existing_nullable=False)
|
||||
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade():
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
with op.batch_alter_table('tenant_preferred_model_providers', schema=None) as batch_op:
|
||||
batch_op.alter_column('provider_name',
|
||||
existing_type=sa.String(length=255),
|
||||
type_=sa.VARCHAR(length=40),
|
||||
existing_nullable=False)
|
||||
|
||||
with op.batch_alter_table('tenant_default_models', schema=None) as batch_op:
|
||||
batch_op.alter_column('provider_name',
|
||||
existing_type=sa.String(length=255),
|
||||
type_=sa.VARCHAR(length=40),
|
||||
existing_nullable=False)
|
||||
|
||||
with op.batch_alter_table('providers', schema=None) as batch_op:
|
||||
batch_op.alter_column('provider_name',
|
||||
existing_type=sa.String(length=255),
|
||||
type_=sa.VARCHAR(length=40),
|
||||
existing_nullable=False)
|
||||
|
||||
with op.batch_alter_table('provider_orders', schema=None) as batch_op:
|
||||
batch_op.alter_column('provider_name',
|
||||
existing_type=sa.String(length=255),
|
||||
type_=sa.VARCHAR(length=40),
|
||||
existing_nullable=False)
|
||||
|
||||
with op.batch_alter_table('provider_models', schema=None) as batch_op:
|
||||
batch_op.alter_column('provider_name',
|
||||
existing_type=sa.String(length=255),
|
||||
type_=sa.VARCHAR(length=40),
|
||||
existing_nullable=False)
|
||||
|
||||
with op.batch_alter_table('provider_model_settings', schema=None) as batch_op:
|
||||
batch_op.drop_index('provider_model_setting_tenant_provider_model_idx')
|
||||
|
||||
op.drop_table('provider_model_settings')
|
||||
with op.batch_alter_table('load_balancing_model_configs', schema=None) as batch_op:
|
||||
batch_op.drop_index('load_balancing_model_config_tenant_provider_model_idx')
|
||||
|
||||
op.drop_table('load_balancing_model_configs')
|
||||
# ### end Alembic commands ###
|
||||
@ -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()
|
||||
@ -0,0 +1,77 @@
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from core.entities.provider_entities import ModelLoadBalancingConfiguration
|
||||
from core.model_manager import LBModelManager
|
||||
from core.model_runtime.entities.model_entities import ModelType
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def lb_model_manager():
|
||||
load_balancing_configs = [
|
||||
ModelLoadBalancingConfiguration(
|
||||
id='id1',
|
||||
name='__inherit__',
|
||||
credentials={}
|
||||
),
|
||||
ModelLoadBalancingConfiguration(
|
||||
id='id2',
|
||||
name='first',
|
||||
credentials={"openai_api_key": "fake_key"}
|
||||
),
|
||||
ModelLoadBalancingConfiguration(
|
||||
id='id3',
|
||||
name='second',
|
||||
credentials={"openai_api_key": "fake_key"}
|
||||
)
|
||||
]
|
||||
|
||||
lb_model_manager = LBModelManager(
|
||||
tenant_id='tenant_id',
|
||||
provider='openai',
|
||||
model_type=ModelType.LLM,
|
||||
model='gpt-4',
|
||||
load_balancing_configs=load_balancing_configs,
|
||||
managed_credentials={"openai_api_key": "fake_key"}
|
||||
)
|
||||
|
||||
lb_model_manager.cooldown = MagicMock(return_value=None)
|
||||
|
||||
def is_cooldown(config: ModelLoadBalancingConfiguration):
|
||||
if config.id == 'id1':
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
lb_model_manager.in_cooldown = MagicMock(side_effect=is_cooldown)
|
||||
|
||||
return lb_model_manager
|
||||
|
||||
|
||||
def test_lb_model_manager_fetch_next(mocker, lb_model_manager):
|
||||
assert len(lb_model_manager._load_balancing_configs) == 3
|
||||
|
||||
config1 = lb_model_manager._load_balancing_configs[0]
|
||||
config2 = lb_model_manager._load_balancing_configs[1]
|
||||
config3 = lb_model_manager._load_balancing_configs[2]
|
||||
|
||||
assert lb_model_manager.in_cooldown(config1) is True
|
||||
assert lb_model_manager.in_cooldown(config2) is False
|
||||
assert lb_model_manager.in_cooldown(config3) is False
|
||||
|
||||
start_index = 0
|
||||
def incr(key):
|
||||
nonlocal start_index
|
||||
start_index += 1
|
||||
return start_index
|
||||
|
||||
mocker.patch('redis.Redis.incr', side_effect=incr)
|
||||
mocker.patch('redis.Redis.set', return_value=None)
|
||||
mocker.patch('redis.Redis.expire', return_value=None)
|
||||
|
||||
config = lb_model_manager.fetch_next()
|
||||
assert config == config2
|
||||
|
||||
config = lb_model_manager.fetch_next()
|
||||
assert config == config3
|
||||
@ -0,0 +1,183 @@
|
||||
from core.entities.provider_entities import ModelSettings
|
||||
from core.model_runtime.entities.model_entities import ModelType
|
||||
from core.model_runtime.model_providers import model_provider_factory
|
||||
from core.provider_manager import ProviderManager
|
||||
from models.provider import LoadBalancingModelConfig, ProviderModelSetting
|
||||
|
||||
|
||||
def test__to_model_settings(mocker):
|
||||
# Get all provider entities
|
||||
provider_entities = model_provider_factory.get_providers()
|
||||
|
||||
provider_entity = None
|
||||
for provider in provider_entities:
|
||||
if provider.provider == 'openai':
|
||||
provider_entity = provider
|
||||
|
||||
# Mocking the inputs
|
||||
provider_model_settings = [ProviderModelSetting(
|
||||
id='id',
|
||||
tenant_id='tenant_id',
|
||||
provider_name='openai',
|
||||
model_name='gpt-4',
|
||||
model_type='text-generation',
|
||||
enabled=True,
|
||||
load_balancing_enabled=True
|
||||
)]
|
||||
load_balancing_model_configs = [
|
||||
LoadBalancingModelConfig(
|
||||
id='id1',
|
||||
tenant_id='tenant_id',
|
||||
provider_name='openai',
|
||||
model_name='gpt-4',
|
||||
model_type='text-generation',
|
||||
name='__inherit__',
|
||||
encrypted_config=None,
|
||||
enabled=True
|
||||
),
|
||||
LoadBalancingModelConfig(
|
||||
id='id2',
|
||||
tenant_id='tenant_id',
|
||||
provider_name='openai',
|
||||
model_name='gpt-4',
|
||||
model_type='text-generation',
|
||||
name='first',
|
||||
encrypted_config='{"openai_api_key": "fake_key"}',
|
||||
enabled=True
|
||||
)
|
||||
]
|
||||
|
||||
mocker.patch('core.helper.model_provider_cache.ProviderCredentialsCache.get', return_value={"openai_api_key": "fake_key"})
|
||||
|
||||
provider_manager = ProviderManager()
|
||||
|
||||
# Running the method
|
||||
result = provider_manager._to_model_settings(
|
||||
provider_entity,
|
||||
provider_model_settings,
|
||||
load_balancing_model_configs
|
||||
)
|
||||
|
||||
# Asserting that the result is as expected
|
||||
assert len(result) == 1
|
||||
assert isinstance(result[0], ModelSettings)
|
||||
assert result[0].model == 'gpt-4'
|
||||
assert result[0].model_type == ModelType.LLM
|
||||
assert result[0].enabled is True
|
||||
assert len(result[0].load_balancing_configs) == 2
|
||||
assert result[0].load_balancing_configs[0].name == '__inherit__'
|
||||
assert result[0].load_balancing_configs[1].name == 'first'
|
||||
|
||||
|
||||
def test__to_model_settings_only_one_lb(mocker):
|
||||
# Get all provider entities
|
||||
provider_entities = model_provider_factory.get_providers()
|
||||
|
||||
provider_entity = None
|
||||
for provider in provider_entities:
|
||||
if provider.provider == 'openai':
|
||||
provider_entity = provider
|
||||
|
||||
# Mocking the inputs
|
||||
provider_model_settings = [ProviderModelSetting(
|
||||
id='id',
|
||||
tenant_id='tenant_id',
|
||||
provider_name='openai',
|
||||
model_name='gpt-4',
|
||||
model_type='text-generation',
|
||||
enabled=True,
|
||||
load_balancing_enabled=True
|
||||
)]
|
||||
load_balancing_model_configs = [
|
||||
LoadBalancingModelConfig(
|
||||
id='id1',
|
||||
tenant_id='tenant_id',
|
||||
provider_name='openai',
|
||||
model_name='gpt-4',
|
||||
model_type='text-generation',
|
||||
name='__inherit__',
|
||||
encrypted_config=None,
|
||||
enabled=True
|
||||
)
|
||||
]
|
||||
|
||||
mocker.patch('core.helper.model_provider_cache.ProviderCredentialsCache.get', return_value={"openai_api_key": "fake_key"})
|
||||
|
||||
provider_manager = ProviderManager()
|
||||
|
||||
# Running the method
|
||||
result = provider_manager._to_model_settings(
|
||||
provider_entity,
|
||||
provider_model_settings,
|
||||
load_balancing_model_configs
|
||||
)
|
||||
|
||||
# Asserting that the result is as expected
|
||||
assert len(result) == 1
|
||||
assert isinstance(result[0], ModelSettings)
|
||||
assert result[0].model == 'gpt-4'
|
||||
assert result[0].model_type == ModelType.LLM
|
||||
assert result[0].enabled is True
|
||||
assert len(result[0].load_balancing_configs) == 0
|
||||
|
||||
|
||||
def test__to_model_settings_lb_disabled(mocker):
|
||||
# Get all provider entities
|
||||
provider_entities = model_provider_factory.get_providers()
|
||||
|
||||
provider_entity = None
|
||||
for provider in provider_entities:
|
||||
if provider.provider == 'openai':
|
||||
provider_entity = provider
|
||||
|
||||
# Mocking the inputs
|
||||
provider_model_settings = [ProviderModelSetting(
|
||||
id='id',
|
||||
tenant_id='tenant_id',
|
||||
provider_name='openai',
|
||||
model_name='gpt-4',
|
||||
model_type='text-generation',
|
||||
enabled=True,
|
||||
load_balancing_enabled=False
|
||||
)]
|
||||
load_balancing_model_configs = [
|
||||
LoadBalancingModelConfig(
|
||||
id='id1',
|
||||
tenant_id='tenant_id',
|
||||
provider_name='openai',
|
||||
model_name='gpt-4',
|
||||
model_type='text-generation',
|
||||
name='__inherit__',
|
||||
encrypted_config=None,
|
||||
enabled=True
|
||||
),
|
||||
LoadBalancingModelConfig(
|
||||
id='id2',
|
||||
tenant_id='tenant_id',
|
||||
provider_name='openai',
|
||||
model_name='gpt-4',
|
||||
model_type='text-generation',
|
||||
name='first',
|
||||
encrypted_config='{"openai_api_key": "fake_key"}',
|
||||
enabled=True
|
||||
)
|
||||
]
|
||||
|
||||
mocker.patch('core.helper.model_provider_cache.ProviderCredentialsCache.get', return_value={"openai_api_key": "fake_key"})
|
||||
|
||||
provider_manager = ProviderManager()
|
||||
|
||||
# Running the method
|
||||
result = provider_manager._to_model_settings(
|
||||
provider_entity,
|
||||
provider_model_settings,
|
||||
load_balancing_model_configs
|
||||
)
|
||||
|
||||
# Asserting that the result is as expected
|
||||
assert len(result) == 1
|
||||
assert isinstance(result[0], ModelSettings)
|
||||
assert result[0].model == 'gpt-4'
|
||||
assert result[0].model_type == ModelType.LLM
|
||||
assert result[0].enabled is True
|
||||
assert len(result[0].load_balancing_configs) == 0
|
||||
Loading…
Reference in New Issue