Dify知识库-RAG流程解析

Dify知识库RAG代码流程图

源码解析

document_indexing_task

代码目录：dify/api/tasks/document_indexing_task.py

主要做了以下两件事

1.查询dataset的文章限制是否超出限制，超出抛出异常，将所有document_ids状态改为 error 如果正 常，则更新所有文章状态为 “解析中” parsing

@shared_task(queue='dataset')
def document_indexing_task(dataset_id: str, document_ids: list):
    """
    异步处理文档索引任务。
    :param dataset_id:数据集ID
    :param document_ids:需要处理的文档ID列表

    Usage: document_indexing_task.delay(dataset_id, document_id)
    主要功能：1.查询dataset的文章限制是否超出 入宫超出抛出异常，将所有document_ids状态改为 error 如果正常，则更新所有文章状态为 “解析中” parsing
    2.IndexingRunner.run()中包含了RAG索引的实现细节
    """


    # 初始化文档列表和开始时间
    documents = []
    start_at = time.perf_counter()
    # 从数据库中获取数据集信息
    dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first()

    # 检查文档数量限制
    features = FeatureService.get_features(dataset.tenant_id)
    try:
        if features.billing.enabled:
            # 获取向量空间信息
            vector_space = features.vector_space
            count = len(document_ids)
            # 批量上传限制
            batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT)
            if count > batch_upload_limit:
                raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")

            # 检查是否超过订阅限制
            if 0 < vector_space.limit <= vector_space.size:
                raise ValueError("Your total number of documents plus the number of uploads have over the limit of "
                                 "your subscription.")
    except Exception as e:
        # 如果有异常，更新所有相关文档的状态为错误，并记录异常信息
        for document_id in document_ids:
            document = db.session.query(Document).filter(
                Document.id == document_id,
                Document.dataset_id == dataset_id
            ).first()
            if document:
                document.indexing_status = 'error'
                document.error = str(e)
                document.stopped_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
                db.session.add(document)
        db.session.commit()
        return
    # 更新文档状态为解析中，并添加到处理列表
    for document_id in document_ids:
        logging.info(click.style('Start process document: {}'.format(document_id), fg='green'))

        document = db.session.query(Document).filter(
            Document.id == document_id,
            Document.dataset_id == dataset_id
        ).first()

        if document:
            document.indexing_status = 'parsing'
            document.processing_started_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
            documents.append(document)
            db.session.add(document)
    db.session.commit()
    #......

2.IndexingRunner.run()中包含了RAG索引的实现细节

@shared_task(queue='dataset')
def document_indexing_task(dataset_id: str, document_ids: list):
    """
    异步处理文档索引任务。
    :param dataset_id:数据集ID
    :param document_ids:需要处理的文档ID列表

    Usage: document_indexing_task.delay(dataset_id, document_id)
    主要功能：1.查询dataset的文章限制是否超出 入宫超出抛出异常，将所有document_ids状态改为 error 如果正常，则更新所有文章状态为 “解析中” parsing
    2.IndexingRunner.run()中包含了RAG索引的实现细节
    """
    #......
    # 尝试运行索引处理
    try:
        # RAG索引的实现细节
        indexing_runner = IndexingRunner()
        indexing_runner.run(documents)
        end_at = time.perf_counter()
        logging.info(click.style('Processed dataset: {} latency: {}'.format(dataset_id, end_at - start_at), fg='green'))
    except DocumentIsPausedException as ex:
        # 如果文档被暂停，记录信息
        logging.info(click.style(str(ex), fg='yellow'))
    except Exception:
        pass

IndexingRunner

代码目录：dify/api/core/indexing_runner.py

run() 该函数主要做了以下三件事

1.提取文本

2.转化数据切片

3.将切片后的文本 构造 document_segment 入库

4.索引中间件加载

def run(self, dataset_documents: list[DatasetDocument]):
    """Run the indexing process."""
    """
    运行索引过程，对每个提供的数据集文档进行处理。
    """
    for dataset_document in dataset_documents:
        try:
            # get dataset
            dataset = Dataset.query.filter_by(
                id=dataset_document.dataset_id
            ).first()

            if not dataset:
                raise ValueError("no dataset found")

            # 获取处理规则
            processing_rule = db.session.query(DatasetProcessRule). \
                filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id). \
                first()
            index_type = dataset_document.doc_form  # 文档的形式，用于确定索引处理器类型
            index_processor = IndexProcessorFactory(index_type).init_index_processor()  # 创建索引处理器实例
            # 提取文本数据
            text_docs = self._extract(index_processor, dataset_document, processing_rule.to_dict())
            # print('提取文本数据', text_docs)
            # 转换数据
            documents = self._transform(index_processor, dataset, text_docs, dataset_document.doc_language,
                                        processing_rule.to_dict())
            # print('转换数据', documents)
            # 保存片段 将最终切片后的 chunks 构造 document_segment 入库
            self._load_segments(dataset, dataset_document, documents)

            # load
            self._load(
                index_processor=index_processor,
                dataset=dataset,
                dataset_document=dataset_document,