Underground-Digital/Workflow-Engine

1import datetime2import logging3import time4from typing import Optional5 6import click7from celery import shared_task8from werkzeug.exceptions import NotFound9 10from core.rag.index_processor.index_processor_factory import IndexProcessorFactory11from core.rag.models.document import Document12from extensions.ext_database import db13from extensions.ext_redis import redis_client14from models.dataset import DocumentSegment15 16 17@shared_task(queue="dataset")18def create_segment_to_index_task(segment_id: str, keywords: Optional[list[str]] = None):19    """20    Async create segment to index21    :param segment_id:22    :param keywords:23    Usage: create_segment_to_index_task.delay(segment_id)24    """25    logging.info(click.style("Start create segment to index: {}".format(segment_id), fg="green"))26    start_at = time.perf_counter()27 28    segment = db.session.query(DocumentSegment).filter(DocumentSegment.id == segment_id).first()29    if not segment:30        raise NotFound("Segment not found")31 32    if segment.status != "waiting":33        return34 35    indexing_cache_key = "segment_{}_indexing".format(segment.id)36 37    try:38        # update segment status to indexing39        update_params = {40            DocumentSegment.status: "indexing",41            DocumentSegment.indexing_at: datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),42        }43        DocumentSegment.query.filter_by(id=segment.id).update(update_params)44        db.session.commit()45        document = Document(46            page_content=segment.content,47            metadata={48                "doc_id": segment.index_node_id,49                "doc_hash": segment.index_node_hash,50                "document_id": segment.document_id,51                "dataset_id": segment.dataset_id,52            },53        )54 55        dataset = segment.dataset56 57        if not dataset:58            logging.info(click.style("Segment {} has no dataset, pass.".format(segment.id), fg="cyan"))59            return60 61        dataset_document = segment.document62 63        if not dataset_document:64            logging.info(click.style("Segment {} has no document, pass.".format(segment.id), fg="cyan"))65            return66 67        if not dataset_document.enabled or dataset_document.archived or dataset_document.indexing_status != "completed":68            logging.info(click.style("Segment {} document status is invalid, pass.".format(segment.id), fg="cyan"))69            return70 71        index_type = dataset.doc_form72        index_processor = IndexProcessorFactory(index_type).init_index_processor()73        index_processor.load(dataset, [document])74 75        # update segment to completed76        update_params = {77            DocumentSegment.status: "completed",78            DocumentSegment.completed_at: datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),79        }80        DocumentSegment.query.filter_by(id=segment.id).update(update_params)81        db.session.commit()82 83        end_at = time.perf_counter()84        logging.info(85            click.style("Segment created to index: {} latency: {}".format(segment.id, end_at - start_at), fg="green")86        )87    except Exception as e:88        logging.exception("create segment to index failed")89        segment.enabled = False90        segment.disabled_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)91        segment.status = "error"92        segment.error = str(e)93        db.session.commit()94    finally:95        redis_client.delete(indexing_cache_key)96