当前位置：首页 > article >正文

dify实现原理分析-rag-检索(Retrieval)服务的实现

article 2025/1/31 19:32:57

概述

本文对dify的检索服务的检索过程的实现逻辑进行了分析。通过本文可以对检索服务的检索过程有一个比较清晰的理解，若是要关注实现细节，可以阅读对应部分的代码。

检索的分类

dify实现了三种类型的检索

语义搜索
全文搜索
混合搜索

检索的类型定义如下：

class RetrievalMethod(Enum):
    SEMANTIC_SEARCH = "semantic_search"
    FULL_TEXT_SEARCH = "full_text_search"
    HYBRID_SEARCH = "hybrid_search"

在进行检索的过程中，会根据检索类型来执行不同的流程。

检索服务的总流程

检索服务的实际执行过程在RetrievalService.retrieve(…)函数中实现。该函数的主要实现逻辑如下：

若检索方法为关键词检索（retrieval_method == "keyword_search"），启动关键词检索线程，来根据查询语句来查询相关文档。结果保存到all_documents中。
若检索服务支持语义检索（is_support_semantic_search），则根据参数query查询嵌入向量，返回top_k个最相似的documents。结果保存到all_documents中。
若检索服务支持全文检索（is_support_fulltext_search），则根据参数query在对应向量数据库中进行全文检索。注意，有些向量数据库不支持全文检索。结果保存到all_documents中。
若检索方法是混合检索retrieval_method == RetrievalMethod.HYBRID_SEARCH.value，则启动数据后处理相关流程。结果保存到all_documents中。

关键词检索的实现

关键词检索功能是在函数：RetrievalService.keyword_search()中实现，该函数在 Flask应用程序的上下文中进行关键词搜索操作。该函数的实现逻辑如下：

在给定的应用程序上下文中，从数据库中获取指定ID的数据集。
使用该数据集创建一个关键词检索对象（Keyword），并执行关键词搜索。检索的时候，会根据关键词的存储类来构建检索的类。默认是：JIEBA。然后再调用Keyword对象的search函数来查询对应关键词的文档。
将找到的文档添加到传入的 all_documents 列表中。
如果发生异常，将错误信息添加到 exceptions 列表中。

关键词搜索search函数的执行逻辑

该函数从关键词表中查询关键词列表，并去掉停用词，然后对关键词所在的文档id进行计数和排序。具体的实现逻辑如下：

获取保存数据集关键词的表名，若该数据表不存在，则创建一个。数据表名为：dataset_keyword_tables。
获取top_k参数值，默认为4；
在数据集关键词表中，根据查询语句获取关键词列表，对每个关键词对应的文本ID进行计数，然后根据计数值对文本片段(chunk)ID进行从大到小排序。
从DocumentSegment表中查出对应的segment的详细内容，然后把segment的内容保存到Document对象中，返回Document列表。

语义检索(semantic_search)的实现

若检索方法支持语义检索RetrievalMethod.is_support_semantic_search(retrieval_method)，则会进行启动语义检索线程。该线程调用RetrievalService.embedding_search函数，并把结果保存到all_documents结果列表中。

让我们来看一下RetrievalService.embedding_search函数的实现逻辑：

查询数据集id对应的数据集；
在向量数据库种搜素与查询文档相似的嵌入向量，查询过程可以参考《dify实现原理分析-rag-文本的嵌入向量的计算和存储第2步》的“嵌入向量的查询”章节的内容。这里不再赘述。

全文检索(full_text_search)的实现

若检索方法支持全文检索RetrievalMethod.is_support_fulltext_search(retrieval_method)，则会进行启动全文检索线程。该线程调用RetrievalService.full_text_index_search函数，并把结果保存到all_documents结果列表中。

查询对应数据集id的数据集，然后创建一个向量数据库对象。vector = Vector(dataset)。
调用vector.search_by_full_text函数通过全文检索，找到最相似的文档。不同向量数据库实现的全文检索的方式不同。
全文检索的过程是通过bm25算法来进行搜索的。有些向量数据库不支持bm25算法，也就不支持全文检索方式。
这里通过用PGVecotor举例说明全文检索的具体实现：

在PGVector中执行全文检索，其实就是执行了一条SQL语句，该语句如下：

                f"""SELECT meta, text, ts_rank(to_tsvector(coalesce(text, '')), plainto_tsquery(%s)) AS score
                FROM {self.table_name}
                WHERE to_tsvector(text) @@ plainto_tsquery(%s)
                ORDER BY score DESC
                LIMIT {top_k}""",
                # f"'{query}'" is required in order to account for whitespace in query
                (f"'{query}'", f"'{query}'"),
            )

其中query是给出的查询内容。

混合检索(hybrid_search)的实现

若retrieval_method的方法是混合检索，执行的过程如下：

        # 混合搜索
        if retrieval_method == RetrievalMethod.HYBRID_SEARCH.value:
            data_post_processor = DataPostProcessor(
                str(dataset.tenant_id), reranking_mode, reranking_model, weights, False
            )
            all_documents = data_post_processor.invoke(
                query=query, documents=all_documents, score_threshold=score_threshold, top_n=top_k
            )