Underground-Digital/Workflow-Engine

1import datetime2import logging3import time4 5import click6from celery import shared_task7 8from configs import dify_config9from core.indexing_runner import DocumentIsPausedError, IndexingRunner10from extensions.ext_database import db11from models.dataset import Dataset, Document12from services.feature_service import FeatureService13 14 15@shared_task(queue="dataset")16def document_indexing_task(dataset_id: str, document_ids: list):17    """18    Async process document19    :param dataset_id:20    :param document_ids:21 22    Usage: document_indexing_task.delay(dataset_id, document_id)23    """24    documents = []25    start_at = time.perf_counter()26 27    dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first()28 29    # check document limit30    features = FeatureService.get_features(dataset.tenant_id)31    try:32        if features.billing.enabled:33            vector_space = features.vector_space34            count = len(document_ids)35            batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT)36            if count > batch_upload_limit:37                raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")38            if 0 < vector_space.limit <= vector_space.size:39                raise ValueError(40                    "Your total number of documents plus the number of uploads have over the limit of "41                    "your subscription."42                )43    except Exception as e:44        for document_id in document_ids:45            document = (46                db.session.query(Document).filter(Document.id == document_id, Document.dataset_id == dataset_id).first()47            )48            if document:49                document.indexing_status = "error"50                document.error = str(e)51                document.stopped_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)52                db.session.add(document)53        db.session.commit()54        return55 56    for document_id in document_ids:57        logging.info(click.style("Start process document: {}".format(document_id), fg="green"))58 59        document = (60            db.session.query(Document).filter(Document.id == document_id, Document.dataset_id == dataset_id).first()61        )62 63        if document:64            document.indexing_status = "parsing"65            document.processing_started_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)66            documents.append(document)67            db.session.add(document)68    db.session.commit()69 70    try:71        indexing_runner = IndexingRunner()72        indexing_runner.run(documents)73        end_at = time.perf_counter()74        logging.info(click.style("Processed dataset: {} latency: {}".format(dataset_id, end_at - start_at), fg="green"))75    except DocumentIsPausedError as ex:76        logging.info(click.style(str(ex), fg="yellow"))77    except Exception:78        pass79