RAGflow/api/db/services/task_service.py

#
#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import os
import random
import xxhash
from datetime import datetime

from api.db.db_utils import bulk_insert_into_db
from deepdoc.parser import PdfParser
from peewee import JOIN
from api.db.db_models import DB, File2Document, File
from api.db import StatusEnum, FileType, TaskStatus
from api.db.db_models import Task, Document, Knowledgebase, Tenant
from api.db.services.common_service import CommonService
from api.db.services.document_service import DocumentService
from api.utils import current_timestamp, get_uuid
from deepdoc.parser.excel_parser import RAGFlowExcelParser
from rag.settings import SVR_QUEUE_NAME
from rag.utils.storage_factory import STORAGE_IMPL
from rag.utils.redis_conn import REDIS_CONN
from api import settings
from rag.nlp import search


def trim_header_by_lines(text: str, max_length) -> str:
    len_text = len(text)
    if len_text <= max_length:
        return text
    for i in range(len_text):
        if text[i] == '\n' and len_text - i <= max_length:
            return text[i + 1:]
    return text


class TaskService(CommonService):
    model = Task

    @classmethod
    @DB.connection_context()
    def get_task(cls, task_id):
        fields = [
            cls.model.id,
            cls.model.doc_id,
            cls.model.from_page,
            cls.model.to_page,
            cls.model.retry_count,
            Document.kb_id,
            Document.parser_id,
            Document.parser_config,
            Document.name,
            Document.type,
            Document.location,
            Document.size,
            Knowledgebase.tenant_id,
            Knowledgebase.language,
            Knowledgebase.embd_id,
            Knowledgebase.pagerank,
            Knowledgebase.parser_config.alias("kb_parser_config"),
            Tenant.img2txt_id,
            Tenant.asr_id,
            Tenant.llm_id,
            cls.model.update_time,
        ]
        docs = (
            cls.model.select(*fields)
                .join(Document, on=(cls.model.doc_id == Document.id))
                .join(Knowledgebase, on=(Document.kb_id == Knowledgebase.id))
                .join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id))
                .where(cls.model.id == task_id)
        )
        docs = list(docs.dicts())
        if not docs:
            return None

        msg = f"\n{datetime.now().strftime('%H:%M:%S')} Task has been received."
        prog = random.random() / 10.0
        if docs[0]["retry_count"] >= 3:
            msg = "\nERROR: Task is abandoned after 3 times attempts."
            prog = -1

        cls.model.update(
            progress_msg=cls.model.progress_msg + msg,
            progress=prog,
            retry_count=docs[0]["retry_count"] + 1,
        ).where(cls.model.id == docs[0]["id"]).execute()

        if docs[0]["retry_count"] >= 3:
            return None

        return docs[0]

    @classmethod
    @DB.connection_context()
    def get_tasks(cls, doc_id: str):
        fields = [
            cls.model.id,
            cls.model.from_page,
            cls.model.progress,
            cls.model.digest,
            cls.model.chunk_ids,
        ]
        tasks = (
            cls.model.select(*fields).order_by(cls.model.from_page.asc(), cls.model.create_time.desc())
                .where(cls.model.doc_id == doc_id)
        )
        tasks = list(tasks.dicts())
        if not tasks:
            return None
        return tasks

    @classmethod
    @DB.connection_context()
    def update_chunk_ids(cls, id: str, chunk_ids: str):
        cls.model.update(chunk_ids=chunk_ids).where(cls.model.id == id).execute()

    @classmethod
    @DB.connection_context()
    def get_ongoing_doc_name(cls):
        with DB.lock("get_task", -1):
            docs = (
                cls.model.select(
                    *[Document.id, Document.kb_id, Document.location, File.parent_id]
                )
                    .join(Document, on=(cls.model.doc_id == Document.id))
                    .join(
                    File2Document,
                    on=(File2Document.document_id == Document.id),
                    join_type=JOIN.LEFT_OUTER,
                )
                    .join(
                    File,
                    on=(File2Document.file_id == File.id),
                    join_type=JOIN.LEFT_OUTER,
                )
                    .where(
                    Document.status == StatusEnum.VALID.value,
                    Document.run == TaskStatus.RUNNING.value,
                    ~(Document.type == FileType.VIRTUAL.value),
                    cls.model.progress < 1,
                    cls.model.create_time >= current_timestamp() - 1000 * 600,
                )
            )
            docs = list(docs.dicts())
            if not docs:
                return []

            return list(
                set(
                    [
                        (
                            d["parent_id"] if d["parent_id"] else d["kb_id"],
                            d["location"],
                        )
                        for d in docs
                    ]
                )
            )

    @classmethod
    @DB.connection_context()
    def do_cancel(cls, id):
        task = cls.model.get_by_id(id)
        _, doc = DocumentService.get_by_id(task.doc_id)
        return doc.run == TaskStatus.CANCEL.value or doc.progress < 0

    @classmethod
    @DB.connection_context()
    def update_progress(cls, id, info):
        if os.environ.get("MACOS"):
            if info["progress_msg"]:
                task = cls.model.get_by_id(id)
                progress_msg = trim_header_by_lines(task.progress_msg + "\n" + info["progress_msg"], 3000)
                cls.model.update(progress_msg=progress_msg).where(cls.model.id == id).execute()
            if "progress" in info:
                cls.model.update(progress=info["progress"]).where(
                    cls.model.id == id
                ).execute()
            return

        with DB.lock("update_progress", -1):
            if info["progress_msg"]:
                task = cls.model.get_by_id(id)
                progress_msg = trim_header_by_lines(task.progress_msg + "\n" + info["progress_msg"], 3000)
                cls.model.update(progress_msg=progress_msg).where(cls.model.id == id).execute()
            if "progress" in info:
                cls.model.update(progress=info["progress"]).where(
                    cls.model.id == id
                ).execute()


def queue_tasks(doc: dict, bucket: str, name: str):
    """
    将文档解析任务分割并加入队列处理。
    
    该函数根据文档类型(PDF、表格等)将文档分割成多个子任务，计算任务摘要，
    检查是否可以重用之前的任务结果，并将未完成的任务加入Redis队列进行处理。
    
    参数:
        doc (dict): 文档信息字典，包含id、type、parser_id、parser_config等信息
        bucket (str): 存储桶名称
        name (str): 文件名称
        
    流程:
        1. 根据文档类型(PDF/表格)将文档分割成多个子任务
        2. 为每个任务生成唯一摘要(digest)
        3. 尝试重用之前任务的处理结果
        4. 清理旧任务并更新文档状态
        5. 将新任务批量插入数据库
        6. 将未完成的任务加入Redis队列
    """
    def new_task():
        """
        创建一个新的任务字典，包含基本任务信息。
        
        返回:
            dict: 包含任务ID、文档ID、进度和页面范围的任务字典
        """
        return {"id": get_uuid(), "doc_id": doc["id"], "progress": 0.0, "from_page": 0, "to_page": 100000000}

    # 存储所有解析任务的数组
    parse_task_array = []

    # PDF文档处理逻辑
    if doc["type"] == FileType.PDF.value:
        # 从存储中获取文件内容
        file_bin = STORAGE_IMPL.get(bucket, name)
        # 获取布局识别方式，默认为"DeepDOC"
        do_layout = doc["parser_config"].get("layout_recognize", "DeepDOC")
        # 获取PDF总页数
        pages = PdfParser.total_page_number(doc["name"], file_bin)
        # 获取每个任务处理的页数，默认为12页
        page_size = doc["parser_config"].get("task_page_size", 12)
        # 对于学术论文类型，默认任务页数为22
        if doc["parser_id"] == "paper":
            page_size = doc["parser_config"].get("task_page_size", 22)
        # 对于特定解析器或非DeepDOC布局识别，将整个文档作为一个任务处理
        if doc["parser_id"] in ["one", "knowledge_graph"] or do_layout != "DeepDOC":
            page_size = 10 ** 9
        # 获取需要处理的页面范围，默认为全部页面
        page_ranges = doc["parser_config"].get("pages") or [(1, 10 ** 5)]
        # 根据页面范围和任务页数分割任务
        for s, e in page_ranges:
            # 调整页码（从0开始）
            s -= 1
            s = max(0, s)
            # 确保结束页不超过文档总页数
            e = min(e - 1, pages)
            # 按照任务页数分割任务
            for p in range(s, e, page_size):
                task = new_task()
                task["from_page"] = p
                task["to_page"] = min(p + page_size, e)
                parse_task_array.append(task)

    # 表格文档处理逻辑
    elif doc["parser_id"] == "table":
        # 从存储中获取文件内容
        file_bin = STORAGE_IMPL.get(bucket, name)
        # 获取表格总行数
        rn = RAGFlowExcelParser.row_number(doc["name"], file_bin)
        # 每3000行作为一个任务
        for i in range(0, rn, 3000):
            task = new_task()
            task["from_page"] = i
            task["to_page"] = min(i + 3000, rn)
            parse_task_array.append(task)
    # 其他类型文档，整个文档作为一个任务处理
    else:
        parse_task_array.append(new_task())

    # 获取文档的分块配置
    chunking_config = DocumentService.get_chunking_config(doc["id"])
    # 为每个任务生成唯一摘要(digest)
    for task in parse_task_array:
        # 创建哈希对象
        hasher = xxhash.xxh64()
        # 对分块配置中的每个字段进行哈希
        for field in sorted(chunking_config.keys()):
            if field == "parser_config":
                # 移除不需要参与哈希计算的特定配置项
                for k in ["raptor", "graphrag"]:
                    if k in chunking_config[field]:
                        del chunking_config[field][k]
            # 将配置字段添加到哈希计算中
            hasher.update(str(chunking_config[field]).encode("utf-8"))
        # 将任务特定字段添加到哈希计算中
        for field in ["doc_id", "from_page", "to_page"]:
            hasher.update(str(task.get(field, "")).encode("utf-8"))
        # 生成任务摘要并设置初始进度
        task_digest = hasher.hexdigest()
        task["digest"] = task_digest
        task["progress"] = 0.0

    # 获取文档之前的任务记录
    prev_tasks = TaskService.get_tasks(doc["id"])
    # 记录重用的块数量
    ck_num = 0
    if prev_tasks:
        # 尝试重用之前任务的处理结果
        for task in parse_task_array:
            ck_num += reuse_prev_task_chunks(task, prev_tasks, chunking_config)
        # 删除文档之前的任务记录
        TaskService.filter_delete([Task.doc_id == doc["id"]])
        # 收集需要删除的块ID
        chunk_ids = []
        for task in prev_tasks:
            if task["chunk_ids"]:
                chunk_ids.extend(task["chunk_ids"].split())
        # 从文档存储中删除这些块
        if chunk_ids:
            settings.docStoreConn.delete({"id": chunk_ids}, search.index_name(chunking_config["tenant_id"]),
                                         chunking_config["kb_id"])
    # 更新文档的块数量
    DocumentService.update_by_id(doc["id"], {"chunk_num": ck_num})

    # 将新任务批量插入数据库
    bulk_insert_into_db(Task, parse_task_array, True)
    # 开始解析文档
    DocumentService.begin2parse(doc["id"])

    # 筛选出未完成的任务
    unfinished_task_array = [task for task in parse_task_array if task["progress"] < 1.0]
    # 将未完成的任务加入Redis队列
    for unfinished_task in unfinished_task_array:
        assert REDIS_CONN.queue_product(
            SVR_QUEUE_NAME, message=unfinished_task
        ), "Can't access Redis. Please check the Redis' status."


def reuse_prev_task_chunks(task: dict, prev_tasks: list[dict], chunking_config: dict):
    idx = 0
    while idx < len(prev_tasks):
        prev_task = prev_tasks[idx]
        if prev_task.get("from_page", 0) == task.get("from_page", 0) \
                and prev_task.get("digest", 0) == task.get("digest", ""):
            break
        idx += 1

    if idx >= len(prev_tasks):
        return 0
    prev_task = prev_tasks[idx]
    if prev_task["progress"] < 1.0 or not prev_task["chunk_ids"]:
        return 0
    task["chunk_ids"] = prev_task["chunk_ids"]
    task["progress"] = 1.0
    if "from_page" in task and "to_page" in task and int(task['to_page']) - int(task['from_page']) >= 10 ** 6:
        task["progress_msg"] = f"Page({task['from_page']}~{task['to_page']}): "
    else:
        task["progress_msg"] = ""
    task["progress_msg"] = " ".join(
        [datetime.now().strftime("%H:%M:%S"), task["progress_msg"], "Reused previous task's chunks."])
    prev_task["chunk_ids"] = ""

    return len(task["chunk_ids"].split())