Underground-Digital/Workflow-Engine

1import datetime2import logging3import time4 5import click6from celery import shared_task7from werkzeug.exceptions import NotFound8 9from core.rag.index_processor.index_processor_factory import IndexProcessorFactory10from core.rag.models.document import Document11from extensions.ext_database import db12from extensions.ext_redis import redis_client13from models.dataset import DocumentSegment14 15 16@shared_task(queue="dataset")17def enable_segment_to_index_task(segment_id: str):18    """19    Async enable segment to index20    :param segment_id:21 22    Usage: enable_segment_to_index_task.delay(segment_id)23    """24    logging.info(click.style("Start enable segment to index: {}".format(segment_id), fg="green"))25    start_at = time.perf_counter()26 27    segment = db.session.query(DocumentSegment).filter(DocumentSegment.id == segment_id).first()28    if not segment:29        raise NotFound("Segment not found")30 31    if segment.status != "completed":32        raise NotFound("Segment is not completed, enable action is not allowed.")33 34    indexing_cache_key = "segment_{}_indexing".format(segment.id)35 36    try:37        document = Document(38            page_content=segment.content,39            metadata={40                "doc_id": segment.index_node_id,41                "doc_hash": segment.index_node_hash,42                "document_id": segment.document_id,43                "dataset_id": segment.dataset_id,44            },45        )46 47        dataset = segment.dataset48 49        if not dataset:50            logging.info(click.style("Segment {} has no dataset, pass.".format(segment.id), fg="cyan"))51            return52 53        dataset_document = segment.document54 55        if not dataset_document:56            logging.info(click.style("Segment {} has no document, pass.".format(segment.id), fg="cyan"))57            return58 59        if not dataset_document.enabled or dataset_document.archived or dataset_document.indexing_status != "completed":60            logging.info(click.style("Segment {} document status is invalid, pass.".format(segment.id), fg="cyan"))61            return62 63        index_processor = IndexProcessorFactory(dataset_document.doc_form).init_index_processor()64        # save vector index65        index_processor.load(dataset, [document])66 67        end_at = time.perf_counter()68        logging.info(69            click.style("Segment enabled to index: {} latency: {}".format(segment.id, end_at - start_at), fg="green")70        )71    except Exception as e:72        logging.exception("enable segment to index failed")73        segment.enabled = False74        segment.disabled_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)75        segment.status = "error"76        segment.error = str(e)77        db.session.commit()78    finally:79        redis_client.delete(indexing_cache_key)80