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