dify实现原理分析-rag-数据检索的实现

数据检索的总体执行步骤

数据检索总体步骤如下：

数据检索是在DatasetRetrieval.retrieve函数中实现的，主要实现逻辑分为以下几步：

1. 检查输入的模型、数据集id列表等是否为空；

2. 获取模型实例，并把它转换成LargeLanguageModel对象；

3. 获取模型实例，并获取模型的元数据，主要是模型的各种参数，以及认证参数等。若模型的元数据为空，直接返回None。

4. 默认情况下，规划路由策略被设置为 REACT_ROUTER。如果模型支持工具调用（TOOL_CALL）或多重工具调用（MULTI_TOOL_CALL），则将规划策略更改为 ROUTER；

5. 筛选可用的数据集：若数据集为空，或数据集不可用则过滤掉数据集；后续的数据检索，会从这些可用数据集中来进行检索。

6. 根据配置选择单线程(RetrieveStrategy.SINGLE)或多线程检索(RetrieveStrategy.MULTIPLE)，获取检索到的document列表；这里只是设置检索的参数，而这两种检索方式都会调用检索服务的RetrievalService.retrieve(…)函数来检索符合条件的数据集。

7. 处理dify提供者的document：使用回调函数返回检索结果，根据分数对检索结果进行排序，并返回格式化后的字符串。

   1）获取每个文档的分数（score）的值

   2）查询状态为completed，且可用,doc_id在检索出来的文档列表中的DocumentSegment列表

   3）若segment（文档块）不为空。

   ​ 3.1) 获取segment的id和位置（确定文档内容的读取位置）

   ​ 3.2) 按id所在的position(位置)排序，若id不在字典中排到最后（无穷大inf）

   ​ 3.3) 遍历排好序的segment：根据条件构建新的列表，然后进行一下操作：

   ​ a) 检查每个segment是否包含answer

   ​ b) 包含:构建一个包含问题和答案的字符串

   ​ c) 不包含：则只构建一个问题的字符串

   ​ 3.4) 对已排好序的segment进行遍历

   ​ 3.5) 获取segment对应的dataset_id对应的dataset

   ​ 3.6) 获取segment.document_id对应的document

8. 使用回调函数返回检索结果，根据分数对检索结果进行排序，并返回格式化后的字符串。

检索的详细执行流程

数据检索的函数声明如下：

class DatasetRetrieval:def __init__(self, application_generate_entity=None):self.application_generate_entity = application_generate_entitydef retrieve(self,app_id: str,user_id: str,tenant_id: str,model_config: ModelConfigWithCredentialsEntity,config: DatasetEntity,query: str,invoke_from: InvokeFrom,show_retrieve_source: bool,hit_callback: DatasetIndexToolCallbackHandler,message_id: str,memory: Optional[TokenBufferMemory] = None,) -> Optional[str]:"""Retrieve dataset.:param app_id: app_id:param user_id: user_id:param tenant_id: tenant id:param model_config: model config:param config: dataset config:param query: query:param invoke_from: invoke from:param show_retrieve_source: show retrieve source:param hit_callback: hit callback:param message_id: message id:param memory: memory:return:"""

该函数的详细实现逻辑如下：

1. 检查数据集id列表，若数据集id列表为空，则直接返回None；

        # 检查输入的模型、数据集等是否有效。dataset_ids = config.dataset_idsif len(dataset_ids) == 0:return None

2. 根据模型配置来获取和构建模型实例对象，并获取模型的元数据（各种参数）

        model_type_instance = model_config.provider_model_bundle.model_type_instancemodel_type_instance = cast(LargeLanguageModel, model_type_instance)# 获取模型实例model_manager = ModelManager()model_instance = model_manager.get_model_instance(tenant_id=tenant_id, model_type=ModelType.LLM, provider=model_config.provider, model=model_config.model)

3. 设置计划策略，默认情况下计划策略是：PlanningStrategy.REACT_ROUTER，若模型特征支持工具调用：ModelFeature.TOOL_CALL或MULTI_TOOL_CALL，则把计划策略设置成：ROUTER，即：planning_strategy = PlanningStrategy.ROUTER。

        # 默认情况下，规划策略被设置为 REACT_ROUTER。planning_strategy = PlanningStrategy.REACT_ROUTER# 检查模型的特性（features）。# 如果模型支持工具调用（TOOL_CALL）或多重工具调用（MULTI_TOOL_CALL），则将规划策略更改为 ROUTER。features = model_schema.features# 检查模型是否支持工具调用，若支持计划策略设置为ROUTERif features:if ModelFeature.TOOL_CALL in features or ModelFeature.MULTI_TOOL_CALL in features:planning_strategy = PlanningStrategy.ROUTERavailable_datasets = []

4. 筛选可用数据集：遍历参数中的dataset_ids列表，从数据库中查询对应id的数据集，过滤掉数据集可用文档为0的和数据集的provide为external的数据集。

        # 筛选可用的数据集。for dataset_id in dataset_ids:# 查询对应id列表的数据集dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first()# 数据集为空，pass掉if not dataset:continue# 数据集不可用，pass掉if dataset and dataset.available_document_count == 0 and dataset.provider != "external":continue# 把数据集添加到可用数据集列表中available_datasets.append(dataset)

5. 根据配置选择单线程（single_retrieve）或多线程（multiple_retrieve）检索来检索document，得到结果document列表：all_documents。

        if retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE:all_documents = self.single_retrieve(...)elif retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE:all_documents = self.multiple_retrieve(...)

6. 从all_documents中抽取出：dify_documents（provider == “dify”）和external_documents（provider == “external”）的结果。处理外部和Dify提供者的document，生成相应的上下文和资源信息。

        # 得到不同提供者的documentdify_documents = [item for item in all_documents if item.provider == "dify"]external_documents = [item for item in all_documents if item.provider == "external"]

7. 处理结果队列dify_documents，步骤如下：

（1）收集评分信息（document_score_list）：通过检查 dify_documents 列表中每个文档片段的 score 元数据，构建一个字典 document_score_list，其中键是文档 ID，值是对应的评分。

            # 获取每个文档的分数（score）的值for item in dify_documents:if item.metadata.get("score"):document_score_list[item.metadata["doc_id"]] = item.metadata["score"]

（2）过滤和排序文档片段：根据给定的 dataset_ids 和一些状态条件（如 status=completed, enabled=True），从数据库中查询相关的文档片段（DocumentSegment）。然后将这些文档片段按其在原始列表中的索引顺序进行排序。

（3）构建文档上下文（document_context_list）：对于每个排序后的文档片段，创建一个 DocumentContext 实例，并将其添加到 document_context_list 中。如果文档片段包含答案（answer），则将答案与问题一起作为一个字符串存储在内容字段中；否则，只存储问题。

             # 获取segment的id和位置（确定文档内容的读取位置）index_node_id_to_position = {id: position for position, id in enumerate(index_node_ids)}# 按id所在的position(位置)排序，若id不在字典中排到最后（无穷大inf）sorted_segments = sorted(segments, key=lambda segment: index_node_id_to_position.get(segment.index_node_id, float("inf")))# 遍历排好序的segment：根据条件构建新的列表for segment in sorted_segments:# 检查每个segment是否包含answerif segment.answer: # 包含:构建一个包含问题和答案的字符串document_context_list.append(DocumentContext(content=f"question:{segment.get_sign_content()} answer:{segment.answer}",score=document_score_list.get(segment.index_node_id, None),))else: # 不包含：则只构建一个问题的字符串document_context_list.append(DocumentContext(content=segment.get_sign_content(),score=document_score_list.get(segment.index_node_id, None),))

（4）构建检索资源（retrieval_resource_list）：如果设置了 show_retrieve_source 标志为真，对于每个排序后的文档片段，查询相关的数据集(dataset)和文档(document)信息。创建一个 source 字典，其中包含数据集、文档的详细信息以及文档片段的相关属性（如评分、命中次数、词数等）。将包含详细信息的 source 字典添加到 retrieval_resource_list 中。

 	          if show_retrieve_source: # 设置了展示检索源的标识			for segment in sorted_segments: # 遍历排序segment# 获取segment对应的dataset_id对应的datasetdataset = Dataset.query.filter_by(id=segment.dataset_id).first()# 获取segment.document_id对应的documentdocument = DatasetDocument.query.filter(DatasetDocument.id == segment.document_id,...).first()# 若2者同时存在if dataset and document:# 构建source字典，包含各种信息source = {"dataset_id": dataset.id,"dataset_name": dataset.name,"document_id": document.id,"document_name": document.name,"data_source_type": document.data_source_type,"segment_id": segment.id,"retriever_from": invoke_from.to_source(),"score": document_score_list.get(segment.index_node_id, 0.0),}...# 若segment的回答不为空，则获取：question与answerif segment.answer:source["content"] = f"question:{segment.content} \nanswer:{segment.answer}"else: # 仅获取questionsource["content"] = segment.content# 将源字典添加到retrieval源列表中retrieval_resource_list.append(source)

8. 使用回调函数返回检索结果，根据分数对检索结果进行排序，并返回格式化后的字符串。

        # 使用回调函数返回检索结果，根据分数对检索结果进行排序，并返回格式化后的字符串。                   if hit_callback and retrieval_resource_list:# 根据segment所在doc_id的分数进行排序retrieval_resource_list = sorted(retrieval_resource_list, key=lambda x: x.get("score") or 0.0, reverse=True)# 获取检索列表中的位置参数for position, item in enumerate(retrieval_resource_list, start=1):item["position"] = positionhit_callback.return_retriever_resource_info(retrieval_resource_list)

9. 按分数进行排序，并把文档内容合并在一个字符串中返回

if document_context_list:# 按分数进行排序，并把文档内容合并在一个字符串中返回document_context_list = sorted(document_context_list, key=lambda x: x.score or 0.0, reverse=True)return str("\n".join([document_context.content for document_context in document_context_list]))

总结

总结一下数据检索的主要步骤：（1）参数验证和模型选择；（2）检索策略选择：单线程或多线程检索（3）结果筛选和处理（4）结果合并和格式化处理，然后返回。

不管是单线程检索还是多线程检索，都会调用检索服务的retrieve函数来实现检索功能，检索服务的检索具体实现会在后面的文章中进行分析。