召回测试源码分析

召回测试源码分析

接口：api.add_resource(HitTestingApi, ‘/datasets/uuid:dataset_id/hit-testing’)

源码目录：dify/api/controllers/console/datasets/hit_testing.py

HitTestingApi

HitTestingService 类用于处理针对特定数据集的命中测试(hit testing)请求

1.权限验证/解析校验参数

class HitTestingApi(Resource):

    @setup_required
    @login_required
    @account_initialization_required
    def post(self, dataset_id):
        dataset_id_str = str(dataset_id)
        # 从数据集中获取指定ID的数据集
        dataset = DatasetService.get_dataset(dataset_id_str)
        if dataset is None:
            raise NotFound("Dataset not found.")
        # 检查当前用户是否有访问该数据集的权限
        try:
            DatasetService.check_dataset_permission(dataset, current_user)
        except services.errors.account.NoPermissionError as e:
            raise Forbidden(str(e))
        # 解析请求参数
        parser = reqparse.RequestParser()
        parser.add_argument('query', type=str, location='json') #问题
        parser.add_argument('retrieval_model', type=dict, required=False, location='json') #模型
        args = parser.parse_args()
        # 验证命中测试所需的参数
        HitTestingService.hit_testing_args_check(args)

2.命中测试执行/异常处理：

class HitTestingApi(Resource):

    @setup_required
    @login_required
    @account_initialization_required
    def post(self, dataset_id):
        #......

        try:# 执行命中测试
            response = HitTestingService.retrieve(
                dataset=dataset,
                query=args['query'],
                account=current_user,
                retrieval_model=args['retrieval_model'],
                limit=10
            )
            #构建响应数据
            return {"query": response['query'], 'records': marshal(response['records'], hit_testing_record_fields)}
        #异常处理
        except services.errors.index.IndexNotInitializedError:
            raise DatasetNotInitializedError()
        except ProviderTokenNotInitError as ex:
            raise ProviderNotInitializeError(ex.description)
        except QuotaExceededError:
            raise ProviderQuotaExceededError()
        except ModelCurrentlyNotSupportError:
            raise ProviderModelCurrentlyNotSupportError()
        except LLMBadRequestError:
            raise ProviderNotInitializeError(
                "No Embedding Model or Reranking Model available. Please configure a valid provider "
                "in the Settings -> Model Provider.")
        except InvokeError as e:
            raise CompletionRequestError(e.description)
        except ValueError as e:
            raise ValueError(str(e))
        except Exception as e:
            logging.exception("Hit testing failed.")
            raise InternalServerError(str(e))

HitTestingService.retrieve()

代码目录：api/services/hit_testing_service.py

处理详细的检索逻辑，使用检索服务执行检索操作，创建查询记录

class HitTestingService:
    @classmethod
    def retrieve(cls, dataset: Dataset, query: str, account: Account, retrieval_model: dict, limit: int = 10) -> dict:
        """
        执行命中测试的检索操作。

        :param dataset: 数据集对象。
        :param query: 查询字符串。
        :param account: 用户账户对象。
        :param retrieval_model: 检索模型配置字典。
        :param limit: 返回结果的最大数量，默认为10。
        :return: 包含查询和检索结果的字典。
        """# 如果数据集中可用的文档或段落数量为0，则直接返回空结果
        if dataset.available_document_count == 0 or dataset.available_segment_count == 0:
            return {
                "query": {
                    "content": query,
                    "tsne_position": {'x': 0, 'y': 0},
                },
                "records": []
            }

        start = time.perf_counter()

        #  # 如果没有提供检索模型，则使用数据集默认的检索模型，如果没有则使用全局默认模型
        if not retrieval_model:
            retrieval_model = dataset.retrieval_model if dataset.retrieval_model else default_retrieval_model
        # 使用检索服务执行检索操作
        all_documents = RetrievalService.retrieve(retrival_method=retrieval_model['search_method'],
                                                  dataset_id=dataset.id,
                                                  query=query,
                                                  top_k=retrieval_model['top_k'],
                                                  score_threshold=retrieval_model['score_threshold']
                                                  if retrieval_model['score_threshold_enabled'] else None,
                                                  reranking_model=retrieval_model['reranking_model']
                                                  if retrieval_model['reranking_enable'] else None
                                                  )

        end = time.perf_counter()
        logging.debug(f"Hit testing retrieve in {end - start:0.4f} seconds")
        # 创建数据集查询记录
        dataset_query = DatasetQuery(
            dataset_id=dataset.id,
            content=query,
            source='hit_testing',
            created_by_role='account',
            created_by=account.id
        )
        # 将查询记录添加到数据库会话并提交
        db.session.add(dataset_query)
        db.session.commit()
        # 返回紧凑的检索响应
        return cls.compact_retrieve_response(dataset, query, all_documents)

RetrievalService.retrieve()

代码目录：api/core/rag/datasource/retrieval_service.py

根据指定的检索方法执行检索操作。full_text_search 全局检索 semantic_search 向量检索 keyword_search 关键字检索 hybrid_search 混合检索

class RetrievalService:

    @classmethod
    def retrieve(cls, retrival_method: str, dataset_id: str, query: str,
                 top_k: int, score_threshold: Optional[float] = .0, reranking_model: Optional[dict] = None):
        _"""_
_        根据指定的检索方法执行检索操作。_

_        :param retrival_method: 检索方法类型。 semantic_search full_text_search keyword_search_
_        :param dataset_id: 数据集ID。_
_        :param query: 检索查询字符串。_
_        :param top_k: 返回结果的最大数量。_
_        :param score_threshold: 分数阈值，用于过滤结果。_
_        :param reranking_model: 重排序模型配置，用于对结果进行二次排序。_
_        :return: 检索结果列表。_
_        """_
_        _# print('检索方法类型',retrival_method)
        # 从数据库中查询指定ID的数据集
        dataset = db.session.query(Dataset).filter(
            Dataset.id == dataset_id
        ).first()
        if not dataset or dataset.available_document_count == 0 or dataset.available_segment_count == 0:
            return []
        all_documents = []  # 存储所有检索结果
        threads = [] # 存储执行检索的线程
        exceptions = [] # 存储执行过程中遇到的异常
       _# 关键词检索_
        if retrival_method == 'keyword_search':
            keyword_thread = threading.Thread(target=RetrievalService.keyword_search, kwargs={
                'flask_app': current_app._get_current_object(),
                'dataset_id': dataset_id,
                'query': query,
                'top_k': top_k,
                'all_documents': all_documents,
                'exceptions': exceptions,
            })
            threads.append(keyword_thread)
            keyword_thread.start()
       # 向量检索（混合检索中也会调用）

        if RetrievalMethod.is_support_semantic_search(retrival_method):
            embedding_thread = threading.Thread(target=RetrievalService.embedding_search, kwargs={
                'flask_app': current_app._get_current_object(),
                'dataset_id': dataset_id,
                'query': query,
                'top_k': top_k,
                'score_threshold': score_threshold,
                'reranking_model': reranking_model,
                'all_documents': all_documents,
                'retrival_method': retrival_method,
                'exceptions': exceptions,
            })
            threads.append(embedding_thread)
            embedding_thread.start()

        # 文本检索（混合检索中也会调用）
        if RetrievalMethod.is_support_fulltext_search(retrival_method):
            full_text_index_thread = threading.Thread(target=RetrievalService.full_text_index_search, kwargs={
                'flask_app': current_app._get_current_object(),
                'dataset_id': dataset_id,
                'query': query,
                'retrival_method': retrival_method,
                'score_threshold': score_threshold,
                'top_k': top_k,
                'reranking_model': reranking_model,
                'all_documents': all_documents,
                'exceptions': exceptions,
            })
            threads.append(full_text_index_thread)
            full_text_index_thread.start()
        # 等待所有线程完成
        for thread in threads:
            thread.join()
        # 如果执行过程中有异常，则合并异常信息并抛出
        if exceptions:
            exception_message = ';\n'.join(exceptions)
            raise Exception(exception_message)
       # 混合检索之后会执行向量和文本检索结果合并后的重排序
        if retrival_method == RetrievalMethod.HYBRID_SEARCH.value:
            data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False)
            all_documents = data_post_processor.invoke(
                query=query,
                documents=all_documents,
                score_threshold=score_threshold,
                top_n=top_k
            )
            # print("混合搜索",all_documents)
        return all_documents

RetrievalMethod.is_support_semantic_search 会判断是否是 'semantic_search’或者 ‘hybrid_search’

假设以 semantic_search 为例子，最终会调用 RetrievalService.embedding_search 方法

RetrievalService.embedding_search 该方法会从基础的向量检索，再通过 Rerank 模型对基础数据再进行结果重排序

def embedding_search(cls, flask_app: Flask, dataset_id: str, query: str,
                     top_k: int, score_threshold: Optional[float], reranking_model: Optional[dict],
                     all_documents: list, retrival_method: str, exceptions: list):
    _"""_
_    执行基于向量嵌入的语义搜索。_

_    :param flask_app: Flask应用实例。_
_    :param dataset_id: 数据集ID。_
_    :param query: 搜索查询字符串。_
_    :param top_k: 返回结果的最大数量。_
_    :param score_threshold: 分数阈值，用于过滤结果。_
_    :param reranking_model: 重排序模型配置，用于对结果进行二次排序。_
_    :param all_documents: 存储所有检索结果的列表。_
_    :param retrival_method: 检索方法类型。_
_    :param exceptions: 存储执行过程中遇到的异常的列表。_
_    """_
_    _with flask_app.app_context():
        try:
            # 从数据库中查询指定ID的数据集
            dataset = db.session.query(Dataset).filter(
                Dataset.id == dataset_id
            ).first()
            # 初始化向量检索器
            vector = Vector(
                dataset=dataset
            )
            # 执行向量搜索，根据相似度分数阈值筛选结果
            documents = vector.search_by_vector(
                query,
                search_type='similarity_score_threshold',
                top_k=top_k,
                score_threshold=score_threshold,
                filter={
                    'group_id': [dataset.id]
                }
            )
            # 如果有搜索结果
            if documents:
                # 如果启用了重排序模型且检索方法为语义搜索，则进行结果重排序
                if reranking_model and retrival_method == RetrievalMethod.SEMANTIC_SEARCH.value:
                    data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False)
                    # 调用数据后处理器对结果进行重排序
                    all_documents.extend(data_post_processor.invoke(
                        query=query,
                        documents=documents,
                        score_threshold=score_threshold,
                        top_n=len(documents)
                    ))
                else:
                    # 否则直接将结果添加到总结果列表中
                    all_documents.extend(documents)
        except Exception as e:
            # 如果执行过程中出现异常，则记录异常信息
            exceptions.append(str(e))

DataPostProcessor 为 Rerank 的实现逻辑，调用 rerank 模型做精筛

向量详细过程

# 执行向量搜索，根据相似度分数阈值筛选结果
documents = vector.search_by_vector(
    query,
    search_type='similarity_score_threshold',
    top_k=top_k,
    score_threshold=score_threshold,
    filter={
        'group_id': [dataset.id]
    }
)

最终会在 embed_query 函数中实现

def embed_query(self, text: str) -> list[float]:
    _"""_
_    嵌入查询文本以生成向量表示。_

_    :param text: 要嵌入的文本字符串。_
_    :return: 文本的嵌入向量，表示为浮点数列表。_
_    """_
_    _# 使用文档嵌入缓存或如果不存在则存储
    hash = helper.generate_text_hash(text) # 生成文本的哈希值
    embedding_cache_key = f'{self._model_instance.provider}_{self._model_instance.model}_{hash}'
    # 从Redis中获取嵌入向量
    embedding = redis_client.get(embedding_cache_key)
    if embedding: # 如果存在，则设置过期时间并返回解码后的嵌入向量
        redis_client.expire(embedding_cache_key, 600)
        return list(np.frombuffer(base64.b64decode(embedding), dtype="float"))
    try:
        # 调用模型实例生成文本嵌入
        embedding_result = self._model_instance.invoke_text_embedding(
            texts=[text],
            user=self._user
        )

        embedding_results = embedding_result.embeddings[0] # 获取第一个文本的嵌入结果
        embedding_results = (embedding_results / np.linalg.norm(embedding_results)).tolist() # 归一化并转换为列表
    except Exception as ex:
        raise ex

    try:
        # 将嵌入向量编码为Base64并存储到Redis中
        embedding_vector = np.array(embedding_results)
        vector_bytes = embedding_vector.tobytes()  # 转换为字节
        # Transform to Base64
        encoded_vector = base64.b64encode(vector_bytes)
        # Transform to string
        encoded_str = encoded_vector.decode("utf-8")
        redis_client.setex(embedding_cache_key, 600, encoded_str) # 存储到Redis并设置过期时间

    except IntegrityError:
        db.session.rollback()
    except:
        logging.exception('Failed to add embedding to redis')

    return embedding_results

DataPostProcessor

代码目录：/api/core/rag/data_post_processor

class DataPostProcessor:
    _"""Interface for data post-processing document._
_    """_

_    _def __init__(self, tenant_id: str, reranking_model: dict, reorder_enabled: bool = False):
        self.rerank_runner = self._get_rerank_runner(reranking_model, tenant_id)
        self.reorder_runner = self._get_reorder_runner(reorder_enabled)

    def invoke(self, query: str, documents: list[Document], score_threshold: Optional[float] = None,
               top_n: Optional[int] = None, user: Optional[str] = None) -> list[Document]:
        if self.rerank_runner:
            documents = self.rerank_runner.run(query, documents, score_threshold, top_n, user)

        if self.reorder_runner:
            documents = self.reorder_runner.run(documents)

        return documents

实际的实现在 api/core/model_manager.py 中：

**def** **invoke_rerank**(self, query: str, docs: list[str], score_threshold: Optional[float] **=** None,
                    top_n: Optional[int] **=** None,
                    user: Optional[str] **=** None) \
        **->** RerankResult:
    self.model_type_instance **=** cast(RerankModel, self.model_type_instance)

    _# 轮询调用重排序模型，获得重排序得分，并基于得分进行排序 _
    **return** self._round_robin_invoke(
        function**=**self.model_type_instance.invoke,
        model**=**self.model,
        credentials**=**self.credentials,
        query**=**query,
        docs**=**docs,
        score_threshold**=**score_threshold,
        top_n**=**top_n,
        user**=**user
    )

最终，构造检索输出结果给前端