Dify二次开发，模块API封装（Knowledge_Retriever）

原创已于 2025-05-26 09:38:50 修改 · 1.3k 阅读

7 ·

CC 4.0 BY-SA版权

文章标签：

#python #人工智能

于 2025-05-15 11:58:58 首次发布

Dify二次开发，模块API封装（Knowledge_Retriever）

在这里插入图片描述
图片源自官网[1]

声明

首先声明，测试用的是powershell，每个Windows系统都有，用GPT生成的测试代码很方便。作为记录只是部署和必要参数查询教程。Dify版本为1.3.1，针对workflow和chatflow的知识检索模块。
POST 请求
URL:console/api/workflow/knowledge-retriever/fetch-dataset
github地址：https://github.com/First-Comer/dify_knowledge_retrieve_api

Docker部署

Dify作为当今最流行的开源LLM应用开发平台是很多开发人员的优选，但是官方给出的API相对较少，在不改动源代码的情况下进行二次开发，用最简单的模式进行开发和部署。部署方式和源代码相同，但在部署前，我们新建Dockerfile，本文以源代码为基础镜像，并将本地新增的代码复制到新容器中。
Dockerfile：

FROM langgenius/dify-api:1.3.1

COPY ../api /app/api
COPY ../api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py /app/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py

只COPY /api这部分代码我要封装的模块不能复制过去，没找到什么原因很奇怪，暂且当作是源码问题，干脆专门把这个要用到代码复制过来。

再修改源码部署方式，因为本次的开发只是对/api中的小模块进行API封装，所以我们只需要修改/api部分代码的部署即可。
docker-compose.yaml：

  api:
    build:
      context: ..
      dockerfile: docker/Dockerfile
    # 用我们自己打的镜像标签
    image: my-dify-api-test:1.3.1
    restart: always
    environment:
      # Use the shared environment variables.
      <<: *shared-api-worker-env
      # Startup mode, 'api' starts the API server.
      MODE: api
      SENTRY_DSN: ${API_SENTRY_DSN:-}
      SENTRY_TRACES_SAMPLE_RATE: ${API_SENTRY_TRACES_SAMPLE_RATE:-1.0}
      SENTRY_PROFILES_SAMPLE_RATE: ${API_SENTRY_PROFILES_SAMPLE_RATE:-1.0}
      PLUGIN_REMOTE_INSTALL_HOST: ${EXPOSE_PLUGIN_DEBUGGING_HOST:-localhost}
      PLUGIN_REMOTE_INSTALL_PORT: ${EXPOSE_PLUGIN_DEBUGGING_PORT:-5003}
      PLUGIN_MAX_PACKAGE_SIZE: ${PLUGIN_MAX_PACKAGE_SIZE:-52428800}
      INNER_API_KEY_FOR_PLUGIN: ${PLUGIN_DIFY_INNER_API_KEY:-QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1}
    depends_on:
      db:
        condition: service_healthy
      redis:
        condition: service_started
    volumes:
      # Mount the storage directory to the container, for storing user files.
      - ./volumes/app/storage:/app/api/storage
    networks:
      - ssrf_proxy_network
      - default

部署部分的代码就这两段，做二次开发部署自己的镜像都可以用到。
忘了说目录结构，本次开发的全部修改目录都在这！！！

dify-main/
├── api/                          # 你修改过的 Dify API 源码
│   ├── services/   
│   │    ├── workflow/                 
│   │    │    ├── dataset_retriever.py   # 修改此文件
│   ├── controllers/
│   │    ├── console/
│   │    ├── __init__.py                 # 修改此文件
│   │    │    ├── knowledge/         # 新建此文件包
│   │    │    ├── __init__.py          #新建自带
│   │    │    │    ├── rettiever.py    # 新建此文件
├── docker/
│   ├── Dockerfile                # 新建此文件
│   ├── docker-compose.yml      # 修改此文件

API封装代码

直接给代码了，源自GPT-4o和GPT-o4-mini，DeepSeek也有参与，因为免费的份额用完了…

首先是路由注册：
dify-main\api\controllers\console_init_.py

from .knowledge.retriever import KnowledgeRetrieverApi

api.add_resource(
    KnowledgeRetrieverApi,
    "/workflow/knowledge-retriever/fetch-dataset",
    endpoint="workflow_knowledge_retriever_fetch"
)

还有两个代码一起给了。

# api/services/workflow/dataset_retriever.py

from core.workflow.nodes.knowledge_retrieval.knowledge_retrieval_node import KnowledgeRetrievalNode
from core.workflow.nodes.knowledge_retrieval.entities import KnowledgeRetrievalNodeData
from core.variables import StringSegment

def fetch_dataset_retriever(payload: dict) -> list[dict]:
    # 1. 从 payload 中单独提取 id，剩余字段给 Pydantic
    raw = payload["node_data"].copy()
    node_id = raw.pop("id")
    node_data = KnowledgeRetrievalNodeData(**raw)

    # 2. 提取 query
    selector = node_data.query_variable_selector[0]
    query_str = payload["inputs"][selector]
    if not isinstance(query_str, str):
        raise ValueError("Query must be a string")

    # 3. 初始化 node
    node = KnowledgeRetrievalNode.__new__(KnowledgeRetrievalNode)
    node.id = node_id
    node.node_data = node_data

    # 4. 模拟 runtime state
    class DummyPool:
        def __init__(self, value): self.value = value
        def get(self, _): return StringSegment(self.value)

    class DummyRuntime:
        variable_pool = DummyPool(query_str)

    setattr(node, "graph_runtime_state", DummyRuntime())
    setattr(node, f"graph_{node.id}_runtime_state", DummyRuntime())

    # 5. 补齐其他必需属性
    node.tenant_id = payload["tenant_id"]
    node.user_id = payload["user_id"]
    node.app_id = payload["app_id"]
    node.user_from = type("UserFrom", (), {"value": payload.get("user_from", "api")})()

    # 6. 调用底层检索方法
    results = node._fetch_dataset_retriever(node_data=node_data, query=query_str)

    # 7. 转成 JSON-serializable 的 list[dict]
    out: list[dict] = []
    for r in results:
        try:
            out.append(dict(r))
        except Exception:
            out.append(r)
    return out

# api/controllers/console/knowledge/retriever.py

from flask_restful import Resource, request
from services.workflow.dataset_retriever import fetch_dataset_retriever

class KnowledgeRetrieverApi(Resource):
    def post(self):
        payload = request.get_json(force=True)
        result = fetch_dataset_retriever(payload)
        return {"status": "success", "data": result}, 200

直接CV。。。

这些都CV后就可以部署了，我也CV一下源代码的部署README：

快速启动

启动 Dify 服务器的最简单方法是运行我们的 docker-compose.yml 文件。在运行安装命令之前，请确保您的机器上安装了 Docker 和 Docker Compose：

cd docker
cp .env.example .env
docker compose up -d

运行后，可以在浏览器上访问 http://localhost/install 进入 Dify 控制台并开始初始化安装操作。
ok了，非常方便，后面我们测试还需要一些参数在下一章。

参数规范

先给测试代码，这个前提是部署的Dify有workflow或是chatflow，有上传和embedding好了的知识库。

# 构造请求体
$bodyB = @{
  node_data = @{
    id                      = "kr-node-1"                              #知识检索节点的唯一标识符。通常用于区分不同的节点，无需修改
    title                   = "KnowledgeRetrievalNode"      
    type                    = "knowledge-retrieval"
    query_variable_selector = @("question")
    dataset_ids             = @("*****")   #知识库数据id，可以多个，英文逗号隔开","
    retrieval_mode          = "single"                               #指定检索模式。常见的值可能是 "single" 或 "batch"
    single_retrieval_config = @{
      model = @{
        provider          = "tongyi"                                   #自行换
        name              = "qwen-max"
        mode              = "chat"
        completion_params = @{ temperature = 0.7 }
      }
      top_k           = 5
      score_threshold = 0.1
    }
  }
  inputs    = @{ question = "***？" }
  tenant_id = "****"         #租户 ID，标识不同用户
  app_id    = "****"             #应用程序 ID，用于区分不同的应用
  user_id   = "****"             #用户 ID，表示发起请求的用户
  user_from = "api"
}

# 序列化并发送
$jsonB = $bodyB | ConvertTo-Json -Depth 5
Invoke-RestMethod `
  -Uri 'http://localhost/console/api/workflow/knowledge-retriever/fetch-dataset' `
  -Method POST `
  -ContentType 'application/json; charset=utf-8' `
  -Body $jsonB

你就测吧，这么详细肯定没问题🆗
上文代码给出了必要的参数及其定义，大部分可以不用修改，只有四个是要修改的：

user_id查询
1.1连接到 PostgreSQL 数据库
powershell

docker-compose exec db sh
psql -U postgres

在这里插入图片描述
1.2查看数据库列表
sql

\l  # 列出所有数据库
\c dify  # 连接到 dify 数据库
\dt
SELECT id, email, name FROM accounts LIMIT 5;

在这里插入图片描述

到这目的达到了

tenant_id查询
同user_id查询，最后命令换为SELECT id, name FROM tenants;
app_id查询；

url里面有自己看。
dataset_ids 查询；
同app_id相同，都在url。

ok，大功告成，拿到你自己的id换上去测试就是200。

参考

本地Docker部署：
[1] https://github.com/langgenius/dify