Underground-Digital/Workflow-Engine

1import logging2import time3from typing import Optional4 5import click6from celery import shared_task7 8from core.rag.index_processor.index_processor_factory import IndexProcessorFactory9from extensions.ext_database import db10from extensions.ext_storage import storage11from models.dataset import Dataset, DocumentSegment12from models.model import UploadFile13 14 15@shared_task(queue="dataset")16def clean_document_task(document_id: str, dataset_id: str, doc_form: str, file_id: Optional[str]):17    """18    Clean document when document deleted.19    :param document_id: document id20    :param dataset_id: dataset id21    :param doc_form: doc_form22    :param file_id: file id23 24    Usage: clean_document_task.delay(document_id, dataset_id)25    """26    logging.info(click.style("Start clean document when document deleted: {}".format(document_id), fg="green"))27    start_at = time.perf_counter()28 29    try:30        dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first()31 32        if not dataset:33            raise Exception("Document has no dataset")34 35        segments = db.session.query(DocumentSegment).filter(DocumentSegment.document_id == document_id).all()36        # check segment is exist37        if segments:38            index_node_ids = [segment.index_node_id for segment in segments]39            index_processor = IndexProcessorFactory(doc_form).init_index_processor()40            index_processor.clean(dataset, index_node_ids)41 42            for segment in segments:43                db.session.delete(segment)44 45            db.session.commit()46        if file_id:47            file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first()48            if file:49                try:50                    storage.delete(file.key)51                except Exception:52                    logging.exception("Delete file failed when document deleted, file_id: {}".format(file_id))53                db.session.delete(file)54                db.session.commit()55 56        end_at = time.perf_counter()57        logging.info(58            click.style(59                "Cleaned document when document deleted: {} latency: {}".format(document_id, end_at - start_at),60                fg="green",61            )62        )63    except Exception:64        logging.exception("Cleaned document when document deleted failed")65