Dify智能体平台源码二次开发笔记（4） - 多租户的SAAS版实现

前言

        Dify 的多租户功能是其商业版的标准功能，我们应当尊重其盈利模式。只有保持良性的商业运作，Dify 才能持续发展，并为用户提供更优质的功能。因此，此功能仅限学习使用。
我们的需求是：实现类似 SaaS 版的账号隔离，包括智能体、知识库、插件和模型配置等模块的完全隔离。为此，需要启用多租户功能。

租户创建

添加controller

首先创建api的调用，这个controller代码，我就创建在controllers/console/auth下面，起名：tenant.py

class TenantCreateApi(Resource):"""Resource for creating a new tenant."""@validate_tokendef post(self):"""Create a new tenant."""parser = reqparse.RequestParser()parser.add_argument("name", type=str, required=True, location="json")args = parser.parse_args()# if not FeatureService.get_system_features().is_allow_create_workspace:#     raise NotAllowedCreateWorkspace()tenant = TenantService.create_tenant(args["name"])return {"result": "success", "data": {"tenant_id": str(tenant.id)}}

这里方法很简单，接收一个租户名称，通过调用

TenantService的租户创建方法，把租户创建成功。

这里@validate_token，是我自己集成我自己云中来平台登录验证的装饰器，因为我是放在我云中来平台调用的，不想用dify的登录验证，如果大家需要用dify的登录验证，然后在代码里需要使用

account = current_user这个登录用户来做业务处理，就可以换成：

@setup_required
@login_required
@account_initialization_required.

这三个装饰器，也是没有登录问题的。

配置路由

api.add_resource(TenantCreateApi, "/tenant/create")

import

我是创建在console下面的，记得在__init__.py里，把tenant import 加上

from .auth import activate, data_source_bearer_auth, data_source_oauth, forgot_password, login, oauth , tenant

集成三方系统调用

def validate_token(view: Optional[Callable] = None, *, fetch_user_arg: Optional[FetchUserArg] = None):def decorator(view_func):@wraps(view_func)def decorated_view(*args, **kwargs):auth_header = request.headers.get("Authorization")if auth_header is None or " " not in auth_header:raise Unauthorized("Authorization不通过！！！")auth_scheme, auth_info = auth_header.split(None, 1)auth_scheme = auth_scheme.lower()if auth_scheme != "bearer":raise Unauthorized("Authorization不通过！")if not jolywood_redis_client.exists("authorization:token:" + auth_info):raise Unauthorized("请登录云中来！！")return view_func(*args, **kwargs)return decorated_viewif view is None:return decoratorelse:return decorator(view)

这个比较简单，验证我自己云中来平台的登录token

添加Service

因为dify本身就有

TenantService这个类了，我们就不要再去创建其它的服务类。直接在类中添加我们需要的方法。

@staticmethoddef create_tenant(name: str, is_setup: Optional[bool] = False, is_from_dashboard: Optional[bool] = False) -> Tenant:"""Create tenant"""# if (#     not FeatureService.get_system_features().is_allow_create_workspace#     and not is_setup#     and not is_from_dashboard# ):#     from controllers.console.error import NotAllowedCreateWorkspace##     raise NotAllowedCreateWorkspace()tenant = Tenant(name=name)db.session.add(tenant)db.session.commit()tenant.encrypt_public_key = generate_key_pair(tenant.id)db.session.commit()return tenant

几行代码轻松搞定，dify这里已经写好了密钥的生成方法generate_key_pair，拿来用就行。

结果验证

通过postman的测是调用，租户创建成功

同时我minio中的租户密钥也创建成功

租户查询

和创建一样分成controller和service层的实现

class TenantAllListApi(Resource):@validate_tokendef get(self):parser = reqparse.RequestParser()parser.add_argument("page", type=int, required=False, default=1, location="args")parser.add_argument("limit", type=int, required=False, default=20, location="args")parser.add_argument("search", type=str, required=False, location="args")args = parser.parse_args()tenants = TenantService.get_all_tenants()tenant_list = [{"id": str(tenant.id),"name": tenant.name,"created_at": tenant.created_at.isoformat()} for tenant in tenants]return {"result": "success", "data": tenant_list}

@staticmethoddef get_all_tenants() -> list[Tenant]:"""Get account join tenants"""return (db.session.query(Tenant).all())

这里写的比较简单，没有分页和查询参数，直接把所有租户显示出来

最后

除了创建、查询，接下来更新和删除我就写了，基本大致相差不大。创建完了租户，接下来就是对用户表和用户租户关联表tenant_account_joins的处理，这个下次再说吧。