Dify智能体平台源码二次开发笔记（3） - 智能体API的三方验证集成

天下琴川

于 2025-04-12 14:53:32 发布

阅读量357

点赞数 7

分类专栏： dify 文章标签：笔记人工智能

本文链接：https://blog.youkuaiyun.com/whf102611/article/details/147165197

版权

dify 专栏收录该内容

7 篇文章

订阅专栏

前言

正式进入二次开发阶段，首先说明本次业务场景需求：

我们计划将Dify作为智能体中台服务，前端需要直接调用Dify智能体API。但Dify默认的API验证方式仅需传输API Key即可访问，从安全角度考虑，这种验证方式不适合直接在前端调用，理应在后端进行验证。

然而，为了减少后端传输的复杂度，我们决定改造Dify智能体API的验证规则。目标验证方式为：

Authorization = Bearer 智能体名称 三方平台登录token

核心改造点

1. 改造wraps.py

文件路径：api/controllers/service_api/wraps.py

Dify智能体API通过该文件中的装饰器进行验证和获取智能体模型。我们重点关注@validate_app_token装饰器。

原实现通过validate_and_get_api_token("app")获取智能体密钥，这存在安全风险。我们需要重写该方法以实现三方验证。

重写后的关键代码：

def validate_and_get_api_token(scope: str | None = None):
    """
    Validate and get API token.
    """
    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不通过！")

    # 解析认证信息
    try:
        enum_name, login_token = auth_info.split(" ", 1)
    except ValueError:
        raise Unauthorized("authorization格式不对")

    # 获取对应的 API token
    api_token_value = ApiTokenEnum.get_token(enum_name)
    if not api_token_value:
        raise Unauthorized("api秘钥不正确")

    
    if not jolywood_redis_client.exists("authorization:XXXXX:" + login_token):
        raise Unauthorized("请登录云中来！！")
    

    # 使用枚举对应的 token 值继续原有逻辑
    current_time = datetime.now(UTC).replace(tzinfo=None)
    cutoff_time = current_time - timedelta(minutes=1)

    with Session(db.engine, expire_on_commit=False) as session:
        update_stmt = (
            update(ApiToken)
            .where(
                ApiToken.token == api_token_value,
                (ApiToken.last_used_at.is_(None) | (ApiToken.last_used_at < cutoff_time)),
                ApiToken.type == scope,
            )
            .values(last_used_at=current_time)
            .returning(ApiToken)
        )
        result = session.execute(update_stmt)
        api_token = result.scalar_one_or_none()

        if not api_token:
            stmt = select(ApiToken).where(ApiToken.token == api_token_value, ApiToken.type == scope)
            api_token = session.scalar(stmt)
            if not api_token:
                raise Unauthorized("Access token is invalid")
        else:
            session.commit()

    return api_token

这里取api key的方法，我微调了一下，直接取得写死的枚举方式，但实际也是可以从缓存或者其它方式取得，案例枚举类：

class ApiTokenEnum(Enum):
    LAIBAO = ("laibao", "app-xxxxxxxxxxxxxx")
    SUSHE = ("sushe", "app-xxxxxxxxxxxxxx")
    QINGJIA = ("qingjia", "app-xxxxxxxxxxxxxx")
    CHUQIN = ("chuqin", "app-xxxxxxxxxxxxxx")
    HUIYI = ("huiyi", "app-xxxxxxxxxxxxxx")
    ZHISHIKU = ("zhishiku","dataset-xxxxxxxxxxxxxx")

    def __init__(self, display_name: str, token: str):
        self.display_name = display_name
        self.token = token

    @classmethod
    def get_token(cls, display_name: str) -> str:
        for member in cls:
            if member.display_name == display_name:
                return member.token
        return ""

2. 改造ext_redis.py

文件路径：api/extensions/ext_redis.py

主要新增三方客户端验证逻辑：

jolywood_redis_client.exists("authorization:XXXX:" + login_token)

具体代码实现：

class RedisClientWrapper:
    """
    A wrapper class for the Redis client that addresses the issue where the global
    `redis_client` variable cannot be updated when a new Redis instance is returned
    by Sentinel.

    This class allows for deferred initialization of the Redis client, enabling the
    client to be re-initialized with a new instance when necessary. This is particularly
    useful in scenarios where the Redis instance may change dynamically, such as during
    a failover in a Sentinel-managed Redis setup.

    Attributes:
        _client (redis.Redis): The actual Redis client instance. It remains None until
                               initialized with the `initialize` method.

    Methods:
        initialize(client): Initializes the Redis client if it hasn't been initialized already.
        __getattr__(item): Delegates attribute access to the Redis client, raising an error
                           if the client is not initialized.
    """

    def __init__(self):
        self._client = None

    def initialize(self, client):
        if self._client is None:
            self._client = client

    def __getattr__(self, item):
        if self._client is None:
            raise RuntimeError("Redis client is not initialized. Call init_app first.")
        return getattr(self._client, item)


redis_client = RedisClientWrapper()
# 新增云中来的redis
jolywood_redis_client = RedisClientWrapper() 


def init_app(app: DifyApp):
    # 新增云中来redis变量
    global redis_client, jolywood_redis_client
    connection_class: type[Union[Connection, SSLConnection]] = Connection
    if dify_config.REDIS_USE_SSL:
        connection_class = SSLConnection

    redis_params: dict[str, Any] = {
        "username": dify_config.REDIS_USERNAME,
        "password": dify_config.REDIS_PASSWORD or None,  # Temporary fix for empty password
        "db": dify_config.REDIS_DB,
        "encoding": "utf-8",
        "encoding_errors": "strict",
        "decode_responses": False,
    }

    # 云中来相关 Redis 参数 (db=1)
    jolywood_redis_params = redis_params.copy()
    jolywood_redis_params["db"] = 1

    if dify_config.REDIS_USE_SENTINEL:
        assert dify_config.REDIS_SENTINELS is not None, "REDIS_SENTINELS must be set when REDIS_USE_SENTINEL is True"
        sentinel_hosts = [
            (node.split(":")[0], int(node.split(":")[1])) for node in dify_config.REDIS_SENTINELS.split(",")
        ]
        sentinel = Sentinel(
            sentinel_hosts,
            sentinel_kwargs={
                "socket_timeout": dify_config.REDIS_SENTINEL_SOCKET_TIMEOUT,
                "username": dify_config.REDIS_SENTINEL_USERNAME,
                "password": dify_config.REDIS_SENTINEL_PASSWORD,
            },
        )
        master = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **redis_params)
        redis_client.initialize(master)
        #云中来redis
        jolywood_master = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **jolywood_redis_params)
        jolywood_redis_client.initialize(jolywood_master)

    elif dify_config.REDIS_USE_CLUSTERS:
        assert dify_config.REDIS_CLUSTERS is not None, "REDIS_CLUSTERS must be set when REDIS_USE_CLUSTERS is True"
        nodes = [
            ClusterNode(host=node.split(":")[0], port=int(node.split(":")[1]))
            for node in dify_config.REDIS_CLUSTERS.split(",")
        ]
        # FIXME: mypy error here, try to figure out how to fix it
        redis_client.initialize(RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD))
        #云中来
        jolywood_redis_client.initialize(RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD))
    else:
        redis_params.update(
            {
                "host": dify_config.REDIS_HOST,
                "port": dify_config.REDIS_PORT,
                "connection_class": connection_class,
            }
        )
        pool = redis.ConnectionPool(**redis_params)
        redis_client.initialize(redis.Redis(connection_pool=pool))

        #云中来
        jolywood_redis_params.update({
            "host": dify_config.REDIS_HOST,
            "port": dify_config.REDIS_PORT,
            "connection_class": connection_class,
        })
        jolywood_pool = redis.ConnectionPool(**jolywood_redis_params)
        jolywood_redis_client.initialize(redis.Redis(connection_pool=jolywood_pool))
    app.extensions["redis"] = redis_client
    # 云中来
    app.extensions["jolywood_redis"] = jolywood_redis_client

因为是开发环境，所有我把redis和三方系统集成在一起，dify用的是db=3和4，我三方系统用的就是db 1.