|
|
import logging
|
|
|
import yaml
|
|
|
from typing import Optional
|
|
|
import flask_login
|
|
|
from pathlib import Path
|
|
|
from constants.languages import languages
|
|
|
from core.model_runtime.errors.validate import CredentialsValidateFailedError
|
|
|
from core.plugin.manager.exc import PluginDaemonClientSideError
|
|
|
from extensions.ext_database import db
|
|
|
from models.account import (
|
|
|
Account,
|
|
|
Tenant,
|
|
|
)
|
|
|
from services.account_service import AccountService, TenantService
|
|
|
from services.dataset_service import DatasetService
|
|
|
from services.errors.account import (
|
|
|
AccountRegisterError,
|
|
|
TenantNotFoundError,
|
|
|
)
|
|
|
from services.errors.workspace import WorkSpaceNotAllowedCreateError
|
|
|
from services.ext.dataset_ext_service import DatasetExtService
|
|
|
from services.model_load_balancing_service import ModelLoadBalancingService
|
|
|
from services.model_provider_service import ModelProviderService
|
|
|
from services.plugin.plugin_service import PluginService
|
|
|
from configs import dify_config
|
|
|
from configs.ext_config import get_ext_config
|
|
|
import os
|
|
|
|
|
|
class AccountInfo:
|
|
|
def __init__(self, email, name, user_id, tenant_id):
|
|
|
self.email = email
|
|
|
self.name = name
|
|
|
self.user_id=user_id
|
|
|
self.tenant_id=tenant_id
|
|
|
def to_dict(self):
|
|
|
return {
|
|
|
"tenant_id": self.email,
|
|
|
"tenant_name": self.name,
|
|
|
"api_key": self.tenant_id,
|
|
|
}
|
|
|
|
|
|
class TenantAccountInfo:
|
|
|
def __init__(self, tenant_id:str,
|
|
|
tenant_name:str,
|
|
|
admin_account:str,
|
|
|
admin_account_password:str,
|
|
|
):
|
|
|
self.tenant_id = tenant_id
|
|
|
self.tenant_name = tenant_name
|
|
|
self.admin_account=admin_account
|
|
|
self.admin_account_password=admin_account_password
|
|
|
|
|
|
def to_dict(self):
|
|
|
return {
|
|
|
"tenant_id": self.tenant_id,
|
|
|
"tenant_name": self.tenant_name,
|
|
|
"admin_account": self.admin_account,
|
|
|
"admin_account_password": self.admin_account_password,
|
|
|
}
|
|
|
|
|
|
class TenantData:
|
|
|
def __init__(self,
|
|
|
api_key:str,
|
|
|
dataset_ids:list[str]
|
|
|
):
|
|
|
self.api_key=api_key
|
|
|
self.dataset_ids=dataset_ids
|
|
|
|
|
|
def to_dict(self):
|
|
|
return {
|
|
|
"api_key": self.api_key,
|
|
|
"dataset_ids": self.dataset_ids,
|
|
|
}
|
|
|
|
|
|
class AccountExtService:
|
|
|
|
|
|
@staticmethod
|
|
|
def create_account_and_tenant(
|
|
|
email: str,
|
|
|
name: str,
|
|
|
tenant_name: str,
|
|
|
target_tenant_id: str,
|
|
|
interface_language: Optional[str] = None,
|
|
|
password: Optional[str] = None
|
|
|
) -> Account:
|
|
|
"""create account"""
|
|
|
account = AccountService.create_account(
|
|
|
email=email, name=name, interface_language=interface_language, password=password, is_setup=True
|
|
|
)
|
|
|
account.target_tenant_id = target_tenant_id
|
|
|
TenantService.create_owner_tenant_if_not_exist(account=account,name=tenant_name,is_setup=True)
|
|
|
account.current_tenant.target_tenant_id = target_tenant_id
|
|
|
db.session.commit()
|
|
|
return account
|
|
|
|
|
|
@staticmethod
|
|
|
def get_admin_account() -> Account:
|
|
|
admin = db.session.query(Account).filter(Account.target_tenant_id=="100").first()
|
|
|
return admin
|
|
|
|
|
|
@staticmethod
|
|
|
def update_account_list(
|
|
|
accounts: list[AccountInfo],
|
|
|
target_tenant_id: str,
|
|
|
interface_language: Optional[str] = None,
|
|
|
):
|
|
|
db.session.begin_nested()
|
|
|
"""Register account"""
|
|
|
try:
|
|
|
# 获取对应的企业
|
|
|
tenant = TenantExtService.get_tenant_by_target_tenant_id(target_tenant_id=target_tenant_id)
|
|
|
if tenant is None:
|
|
|
raise TenantNotFoundError("企业未初始，请联系管理员!")
|
|
|
# 获取所有的用户列表
|
|
|
exists = db.session.query(Account).filter(Account.target_tenant_id == target_tenant_id).all()
|
|
|
#
|
|
|
existDict = { account.email: account for account in exists }
|
|
|
|
|
|
for account in accounts:
|
|
|
email = account["email"]
|
|
|
if email in existDict:
|
|
|
existAccount = existDict[email]
|
|
|
existAccount.name = account["name"]
|
|
|
existAccount.email = account["email"]
|
|
|
existAccount.user_id = account["user_id"]
|
|
|
else:
|
|
|
newAccount = AccountService.create_account(email=account["email"],
|
|
|
name=account["name"],
|
|
|
interface_language=interface_language or languages[0],
|
|
|
password="wisdom@123",
|
|
|
is_setup=True)
|
|
|
newAccount.user_id = account["user_id"]
|
|
|
newAccount.target_tenant_id = target_tenant_id
|
|
|
# 创建企业关系
|
|
|
TenantService.create_tenant_member(tenant, newAccount)
|
|
|
db.session.commit()
|
|
|
except WorkSpaceNotAllowedCreateError:
|
|
|
db.session.rollback()
|
|
|
except AccountRegisterError as are:
|
|
|
db.session.rollback()
|
|
|
logging.exception("Register failed")
|
|
|
raise are
|
|
|
except Exception as e:
|
|
|
db.session.rollback()
|
|
|
logging.exception("Register failed")
|
|
|
raise AccountRegisterError(f"Registration failed: {e}") from e
|
|
|
|
|
|
|
|
|
class TenantExtService:
|
|
|
|
|
|
@staticmethod
|
|
|
def get_tenant() -> Tenant:
|
|
|
# 获取第一个企业，为默认企业
|
|
|
tenant = db.session.query(Tenant).first()
|
|
|
return tenant
|
|
|
|
|
|
@staticmethod
|
|
|
def get_tenant_by_target_tenant_id(target_tenant_id:str) -> Tenant:
|
|
|
# 获取第一个企业，为默认企业
|
|
|
tenant = db.session.query(Tenant).filter(Tenant.target_tenant_id == target_tenant_id).first()
|
|
|
return tenant
|
|
|
|
|
|
@staticmethod
|
|
|
def setModeConfig(tenant_id:str, args:dict[str, object], provider:str) -> None:
|
|
|
|
|
|
model_load_balancing_service = ModelLoadBalancingService()
|
|
|
if (
|
|
|
"load_balancing" in args
|
|
|
and args["load_balancing"]
|
|
|
and "enabled" in args["load_balancing"]
|
|
|
and args["load_balancing"]["enabled"]
|
|
|
):
|
|
|
if "configs" not in args["load_balancing"]:
|
|
|
raise ValueError("invalid load balancing configs")
|
|
|
|
|
|
# save load balancing configs
|
|
|
model_load_balancing_service.update_load_balancing_configs(
|
|
|
tenant_id=tenant_id,
|
|
|
provider=provider,
|
|
|
model=args["model"],
|
|
|
model_type=args["model_type"],
|
|
|
configs=args["load_balancing"]["configs"],
|
|
|
)
|
|
|
|
|
|
# enable load balancing
|
|
|
model_load_balancing_service.enable_model_load_balancing(
|
|
|
tenant_id=tenant_id, provider=provider, model=args["model"], model_type=args["model_type"]
|
|
|
)
|
|
|
else:
|
|
|
# disable load balancing
|
|
|
model_load_balancing_service.disable_model_load_balancing(
|
|
|
tenant_id=tenant_id, provider=provider, model=args["model"], model_type=args["model_type"]
|
|
|
)
|
|
|
|
|
|
if args.get("config_from", "") != "predefined-model":
|
|
|
model_provider_service = ModelProviderService()
|
|
|
|
|
|
try:
|
|
|
model_provider_service.save_model_credentials(
|
|
|
tenant_id=tenant_id,
|
|
|
provider=provider,
|
|
|
model=args["model"],
|
|
|
model_type=args["model_type"],
|
|
|
credentials=args["credentials"],
|
|
|
)
|
|
|
except CredentialsValidateFailedError as ex:
|
|
|
logging.exception(
|
|
|
f"Failed to save model credentials, tenant_id: {tenant_id},"
|
|
|
f" model: {args.get('model')}, model_type: {args.get('model_type')}"
|
|
|
)
|
|
|
raise ValueError(str(ex))
|
|
|
|
|
|
@staticmethod
|
|
|
def install_plugin(tenant_id:str):
|
|
|
TenantExtService.install_langgenius(tenant_id=tenant_id)
|
|
|
TenantExtService.install_model(tenant_id=tenant_id)
|
|
|
|
|
|
@staticmethod
|
|
|
def install_model(tenant_id:str):
|
|
|
|
|
|
params = {
|
|
|
"INIT_MODEL_LLM_NAME" : dify_config.INIT_MODEL_LLM_NAME,
|
|
|
"INIT_MODEL_LLM_CONTEXT_SIZE" : dify_config.INIT_MODEL_LLM_CONTEXT_SIZE,
|
|
|
"INIT_MODEL_LLM_MAX_TOKENS" : dify_config.INIT_MODEL_LLM_MAX_TOKENS,
|
|
|
"INIT_MODEL_LLM_BASE_URL" : dify_config.INIT_MODEL_LLM_BASE_URL
|
|
|
}
|
|
|
llm_config = get_ext_config(file_name="plugin_llm_config.yml",params = params)
|
|
|
|
|
|
TenantExtService.setModeConfig(
|
|
|
tenant_id=tenant_id,args=llm_config,provider=dify_config.INIT_MODEL_LLM_PROVIDER
|
|
|
)
|
|
|
|
|
|
params = {
|
|
|
"INIT_MODEL_TEXT_EMBEDDING_NAME" : dify_config.INIT_MODEL_TEXT_EMBEDDING_NAME,
|
|
|
"INIT_MODEL_TEXT_EMBEDDING_CONTEXT_SIZE" : dify_config.INIT_MODEL_TEXT_EMBEDDING_CONTEXT_SIZE,
|
|
|
"INIT_MODEL_TEXT_EMBEDDING_MAX_TOKENS" : dify_config.INIT_MODEL_TEXT_EMBEDDING_MAX_TOKENS,
|
|
|
"INIT_MODEL_TEXT_EMBEDDING_BASE_URL" : dify_config.INIT_MODEL_TEXT_EMBEDDING_BASE_URL
|
|
|
}
|
|
|
text_embedding_config = get_ext_config(file_name="plugin_embedding_config.yml", params=params)
|
|
|
TenantExtService.setModeConfig(
|
|
|
tenant_id=tenant_id,args=text_embedding_config,provider=dify_config.INIT_MODEL_TEXT_EMBEDDING_PROVIDER
|
|
|
)
|
|
|
params = {
|
|
|
"INIT_MODEL_TEXT_EMBEDDING_RERANK_NAME": dify_config.INIT_MODEL_TEXT_EMBEDDING_RERANK_NAME,
|
|
|
"INIT_MODEL_TEXT_EMBEDDING_RERANK_BASE_URL": dify_config.INIT_MODEL_TEXT_EMBEDDING_RERANK_BASE_URL,
|
|
|
}
|
|
|
text_embedding_rerank_config = get_ext_config(file_name="plugin_embedding_rerank_config.yml", params=params)
|
|
|
TenantExtService.setModeConfig(
|
|
|
tenant_id=tenant_id,args=text_embedding_rerank_config,provider=dify_config.INIT_MODEL_TEXT_EMBEDDING_RERANK_PROVIDER
|
|
|
)
|
|
|
|
|
|
@staticmethod
|
|
|
def install_langgenius(tenant_id: str):
|
|
|
upload_unique_identifiers = TenantExtService.upload_langgenius(tenant_id=tenant_id)
|
|
|
# plugin_unique_identifiers = dify_config.PLUGIN_UNIQUE_IDENTIFIERS.split(",") if dify_config.PLUGIN_UNIQUE_IDENTIFIERS else []
|
|
|
# 查询已经安装的
|
|
|
tasks = PluginService.list(tenant_id)
|
|
|
# 已经安装的插件
|
|
|
exists_plugin_unique_identifiers = [item.plugin_unique_identifier for item in tasks]
|
|
|
# 去除已经安装的插件ID，只保留未安装的插件ID
|
|
|
new_unique_identifiers = [uui for uui in upload_unique_identifiers if uui not in exists_plugin_unique_identifiers]
|
|
|
# 安装插件
|
|
|
PluginService.install_from_marketplace_pkg(tenant_id, new_unique_identifiers)
|
|
|
|
|
|
@staticmethod
|
|
|
def upload_langgenius(tenant_id: str) -> list[str]:
|
|
|
directory = Path(__file__).parent.parent.parent / 'plugins' / 'langgenius'
|
|
|
unique_identifiers = []
|
|
|
for filename in os.listdir(directory):
|
|
|
file_path = os.path.join(directory, filename)
|
|
|
if os.path.isfile(file_path):
|
|
|
print(f"读取文件：{file_path}")
|
|
|
with open(file_path, 'rb') as f:
|
|
|
content = f.read()
|
|
|
try:
|
|
|
response = PluginService.upload_pkg(tenant_id=tenant_id, pkg=content)
|
|
|
unique_identifier = response.unique_identifier
|
|
|
unique_identifiers.append(unique_identifier)
|
|
|
except PluginDaemonClientSideError as e:
|
|
|
raise ValueError(e)
|
|
|
return unique_identifiers
|
|
|
|
|
|
@staticmethod
|
|
|
def enable_tenant(
|
|
|
target_tenant_id: str,
|
|
|
target_tenant_name: str,
|
|
|
) -> TenantAccountInfo:
|
|
|
db.session.begin_nested()
|
|
|
password = "wisdom@123"
|
|
|
try:
|
|
|
email = f"admin@{target_tenant_id}.com"
|
|
|
admin_name = f"{target_tenant_name}-管理员"
|
|
|
# 判断企业是否已经创建
|
|
|
tenant = TenantExtService.get_tenant_by_target_tenant_id(target_tenant_id)
|
|
|
if tenant is not None:
|
|
|
account = AccountService.get_user_through_email(email)
|
|
|
if account is None:
|
|
|
account = AccountService.create_account(email=email, name=admin_name, password=password, is_setup=True,interface_language="zh-Hans")
|
|
|
TenantService.create_tenant_member(tenant, account, role="owner")
|
|
|
account.target_tenant_id = target_tenant_id
|
|
|
else:
|
|
|
account = AccountExtService.create_account_and_tenant(email=email,
|
|
|
name=admin_name,
|
|
|
tenant_name=target_tenant_name,
|
|
|
target_tenant_id=target_tenant_id,
|
|
|
interface_language="zh-Hans",
|
|
|
password=password)
|
|
|
# 获取第一个企业，为默认企业
|
|
|
tenant = account.current_tenant
|
|
|
|
|
|
account_info = TenantAccountInfo(tenant_name=tenant.name,
|
|
|
tenant_id=tenant.id,
|
|
|
admin_account=admin_name,
|
|
|
admin_account_password=password)
|
|
|
return account_info
|
|
|
except Exception as e:
|
|
|
db.session.rollback()
|
|
|
logging.exception("Register failed")
|
|
|
raise AccountRegisterError(f"Registration failed: {e}") from e
|
|
|
|
|
|
@staticmethod
|
|
|
def init_tenant(
|
|
|
target_tenant_id: str,
|
|
|
target_tenant_name: str,
|
|
|
) -> TenantData:
|
|
|
db.session.begin_nested()
|
|
|
try:
|
|
|
account = flask_login.current_user
|
|
|
tenant = account.current_tenant
|
|
|
# 初始化大模型插槽
|
|
|
TenantExtService.install_plugin(tenant_id=tenant.id)
|
|
|
# 初始化知识库
|
|
|
datasets = DatasetExtService.init_dataset(
|
|
|
tenant=tenant, target_tenant_id=target_tenant_id,target_tenant_name=target_tenant_name,account=account
|
|
|
)
|
|
|
# 获取Api token
|
|
|
api_token = DatasetExtService().get_or_add_datasets_api_token(tenant_id=tenant.id)
|
|
|
db.session.commit()
|
|
|
dataset_ids = [dataset.id for dataset in datasets]
|
|
|
tenant_data = TenantData(api_key=api_token.token,
|
|
|
dataset_ids=dataset_ids)
|
|
|
return tenant_data
|
|
|
except Exception as e:
|
|
|
db.session.rollback()
|
|
|
logging.exception("Register failed")
|
|
|
raise AccountRegisterError(f"Registration failed: {e}") from e
|
|
|
|
|
|
|