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

概述

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

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

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

Jieba类的构成

关键词索引创建

总体实现逻辑

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

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

该函数的主要逻辑如下：

关键词提取：extract_keywords

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

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

该函数的处理流程如下：

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))

def _expand_tokens_with_subtokens(self, tokens: set[str]) -> set[str]:"""获取tokens的子词，并过滤停用词"""results = set()for token in tokens:# 1. 添加原始tokenresults.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字典中了。

该函数的详细实现如下：

    # 这段代码的主要功能是将一个关键词表 (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_tablekeyword_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"))

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

关键词查询

search函数的详细实现分析：

    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）然后从刚才查询出来的关键词字典中，查询出与查询字符串中关键词匹配的文本索引idsorted_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）把添加完关键词的字典再写回关键词存储表或文件中

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

总结

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