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

前言

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

租户创建

添加controller

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

class TenantCreateApi(Resource):
    """Resource for creating a new tenant."""

    @validate_token
    def 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_view

    if view is None:
        return decorator
    else:
        return decorator(view)

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

添加Service

因为dify本身就有

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

@staticmethod
    def 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_token
    def 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}

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

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

最后

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