dify实现分析-rag-关键词索引的实现

概述

在dify中有两种构建索引的方式，一种是经济型，另一种是高质量索引（通过向量数据库来实现）。其中经济型就是关键词索引，通过构建关键词索引来定位查询的文本块，而关键词索引的构建是通过Jieba这个库来完成的。

Jieba（"结巴"）是一个强大的中文分词和关键词提取工具库。在dify中，Jieba类作为一个基于关键词的文档检索系统的核心实现。

本文介绍关键词索引的构建类Jieba类的实现，包括：文本的索引的添加，修改等操作。

Jieba类的构成

Jieba类基础功能数据管理搜索功能关键词提取分词处理关键词表管理存储策略关键词搜索文档检索

关键词索引创建

总体实现逻辑

关键词索引创建在create函数中实现，该函数的声明如下：

复制代码

    def create(self, texts: list[Document], **kwargs) -> BaseKeyword:

该函数的主要逻辑如下：
开始创建关键词索引获取Redis分布式锁检查关键词表是否存在,若不存在,创建关键词表:dataset_keyword_tables 遍历文件分块列表: Document对象列表关键词提取: extract_keywords 更新段落关键词更新关键词表: 把关键词保存到字典中保存关键词表: 保存关键词表到数据库或文件系统返回实例

关键词提取：extract_keywords

关键词提取主要完成：从文本中提取关键词，支持停用词过滤和子词提取。extract_keywords函数的声明如下：

python 复制代码

    def extract_keywords(self, text: str, max_keywords_per_chunk: Optional[int] = 10) -> set[str]:

该函数的处理流程如下：
输入文本 TFIDF关键词提取子词拆分停用词过滤返回关键词集合

python 复制代码

def extract_keywords(self, text: str, max_keywords_per_chunk: Optional[int] = 10) -> set[str]:
    """Extract keywords with JIEBA tfidf."""
    # 1. 使用TFIDF算法提取关键词
    keywords = jieba.analyse.extract_tags(
        sentence=text,
        topK=max_keywords_per_chunk,  # 默认最多10个关键词
    )

    # 2. 扩展子词并过滤停用词
    return set(self._expand_tokens_with_subtokens(keywords))

python 复制代码

def _expand_tokens_with_subtokens(self, tokens: set[str]) -> set[str]:
    """获取tokens的子词，并过滤停用词"""
    results = set()
    for token in tokens:
        # 1. 添加原始token
        results.add(token)
        # 2. 使用正则提取子词
        sub_tokens = re.findall(r"\w+", token)
        # 3. 如果存在多个子词
        if len(sub_tokens) > 1:
            # 过滤停用词并添加到结果集
            results.update({w for w in sub_tokens if w not in list(STOPWORDS)})
    return results

关键词存储

关键词存储的函数声明如下：

复制代码

def _save_dataset_keyword_table(self, keyword_table):

其中处理完成的关键词，都已经保存到dataset_keyword_tables字典中了。
接收关键词表构建元数据字典判断存储类型数据库存储: database 文件存储: file JSON序列化构建文件路径检查文件是否存在: 若存在删除之保存文件

该函数的详细实现如下：

python 复制代码

    # 这段代码的主要功能是将一个关键词表 (dataset_keyword_tables) 保存到数据库或文件中，具体取决于数据源类型。
    def _save_dataset_keyword_table(self, keyword_table):
        # 创建数据字典，保存元数据信息
        keyword_table_dict = {
            "__type__": "keyword_table",
            "__data__": {"index_id": self.dataset.id, "summary": None, "table": keyword_table},
        }
        # 记录数据集的数据来源类型
        dataset_keyword_table = self.dataset.dataset_keyword_table
        keyword_data_source_type = dataset_keyword_table.data_source_type
        # 数据源是数据库，则将字典编码为 JSON 字符串，并更新数据库中的 keyword_table 字段。然后提交事务。
        if keyword_data_source_type == "database":
            dataset_keyword_table.keyword_table = json.dumps(keyword_table_dict, cls=SetEncoder)
            db.session.commit()
        else:
            # 来源是文件，则构建一个文件键（路径），检查文件是否存在，如果存在则删除
            file_key = "keyword_files/" + self.dataset.tenant_id + "/" + self.dataset.id + ".txt"
            if storage.exists(file_key):
                storage.delete(file_key)
            # 最后将字典编码为 JSON 并保存到指定的文件路径。
            storage.save(file_key, json.dumps(keyword_table_dict, cls=SetEncoder).encode("utf-8"))

python 复制代码

class SetEncoder(json.JSONEncoder):
    """自定义JSON编码器，处理set类型"""
    def default(self, obj):
        if isinstance(obj, set):
            return list(obj)  # 将set转换为list
        return super().default(obj)

关键词查询

获取关键词表名,从表中查询对应数据集的关键词字典获取top-k参数提取查询字符串中的关键词,根据查询关键词在文档切片关键词中进行查询,并对结果排序按排序切片索引id从数据库表中查询文档内容构建Document对象返回Document对象列表

search函数的详细实现分析：

python 复制代码

    def search(self, query: str, **kwargs: Any) -> list[Document]:
        # 从dataset_keyword_tables表中获取对应数据集的数据分块记录字典
        keyword_table = self._get_dataset_keyword_table()

        k = kwargs.get("top_k", 4)

        # （1）使用Jieba对用户输入的查询字符串进行关键词提取
        # （2）然后从刚才查询出来的关键词字典中，查询出与查询字符串中关键词匹配的文本索引id
        sorted_chunk_indices = self._retrieve_ids_by_query(keyword_table, query, k)

        documents = []
        # 根据文本块索引id，从数据库中查询出对应的文本块内容
        for chunk_index in sorted_chunk_indices:
            segment = (
                db.session.query(DocumentSegment)
                .filter(DocumentSegment.dataset_id == self.dataset.id, DocumentSegment.index_node_id == chunk_index)
                .first()
            )

            # 以Document对象结构来返回结果
            if segment:
                documents.append(
                    Document(
                        page_content=segment.content,
                        metadata={
                            "doc_id": chunk_index,
                            "doc_hash": segment.index_node_hash,
                            "document_id": segment.document_id,
                            "dataset_id": segment.dataset_id,
                        },
                    )
                )

        return documents

小结

说明，分析到这里，我们基本上就了解了关键词查询的基本原理：对用户查询字符串进行分词处理（通过Jieba库），根据分词结果在对应数据集中查询对应分词，然后获取该分词对应的文本和文本块。

可见，关键词索引主要是依赖对文本进行分词，然后通过分词来进行匹配，从而找到对应文本块的数据。这种方式并没有从语义角度去理解文本，本质上是通过分词得到的关键词进行匹配的方式来找到对应文本块。与通过语义的方式来查找文本，这种方式会存在一定的局限性。

关键词添加和删除

关键词的添加和删除都是要先从数据表或文件中把该数据集原有关键词读取到一个字典中，然后对该字典中的关键词进行添加或删除操作，然后再把数据写回数据表或文件。

（1）先查询数据集对应的关键词表（或文件）的数据，并以字典的方式返回

（2）在字典中添加对应关键词

（3）把添加完关键词的字典再写回关键词存储表或文件中

关键词的删除和关键词添加步骤类似，只是在第二步会从获取到的字段中把关键词删除，然后再写回数据表或文件中。

总结

关键词索引方式不需要其他额外的存储组件就可以完成索引的构建，成本相对比较低，比较经济实惠。但该方式是通过分词和关键词匹配方式来构建的文本块查询，比起通过向量和语义匹配的方式，有一定的局限性，选择那种方式，需要根据具体的场景来确定。