Underground-Digital/Workflow-Engine

1import datetime2import logging3import time4 5import click6from celery import shared_task7from werkzeug.exceptions import NotFound8 9from core.indexing_runner import DocumentIsPausedError, IndexingRunner10from core.rag.index_processor.index_processor_factory import IndexProcessorFactory11from extensions.ext_database import db12from models.dataset import Dataset, Document, DocumentSegment13 14 15@shared_task(queue="dataset")16def document_indexing_update_task(dataset_id: str, document_id: str):17    """18    Async update document19    :param dataset_id:20    :param document_id:21 22    Usage: document_indexing_update_task.delay(dataset_id, document_id)23    """24    logging.info(click.style("Start update document: {}".format(document_id), fg="green"))25    start_at = time.perf_counter()26 27    document = db.session.query(Document).filter(Document.id == document_id, Document.dataset_id == dataset_id).first()28 29    if not document:30        raise NotFound("Document not found")31 32    document.indexing_status = "parsing"33    document.processing_started_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)34    db.session.commit()35 36    # delete all document segment and index37    try:38        dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first()39        if not dataset:40            raise Exception("Dataset not found")41 42        index_type = document.doc_form43        index_processor = IndexProcessorFactory(index_type).init_index_processor()44 45        segments = db.session.query(DocumentSegment).filter(DocumentSegment.document_id == document_id).all()46        if segments:47            index_node_ids = [segment.index_node_id for segment in segments]48 49            # delete from vector index50            index_processor.clean(dataset, index_node_ids)51 52            for segment in segments:53                db.session.delete(segment)54            db.session.commit()55        end_at = time.perf_counter()56        logging.info(57            click.style(58                "Cleaned document when document update data source or process rule: {} latency: {}".format(59                    document_id, end_at - start_at60                ),61                fg="green",62            )63        )64    except Exception:65        logging.exception("Cleaned document when document update data source or process rule failed")66 67    try:68        indexing_runner = IndexingRunner()69        indexing_runner.run([document])70        end_at = time.perf_counter()71        logging.info(click.style("update document: {} latency: {}".format(document.id, end_at - start_at), fg="green"))72    except DocumentIsPausedError as ex:73        logging.info(click.style(str(ex), fg="yellow"))74    except Exception:75        pass76