Underground-Digital/Workflow-Engine

1import logging2import time3 4import click5from celery import shared_task6from werkzeug.exceptions import NotFound7 8from core.rag.datasource.vdb.vector_factory import Vector9from core.rag.models.document import Document10from extensions.ext_database import db11from extensions.ext_redis import redis_client12from models.dataset import Dataset13from models.model import App, AppAnnotationSetting, MessageAnnotation14from services.dataset_service import DatasetCollectionBindingService15 16 17@shared_task(queue="dataset")18def batch_import_annotations_task(job_id: str, content_list: list[dict], app_id: str, tenant_id: str, user_id: str):19    """20    Add annotation to index.21    :param job_id: job_id22    :param content_list: content list23    :param app_id: app id24    :param tenant_id: tenant id25    :param user_id: user_id26 27    """28    logging.info(click.style("Start batch import annotation: {}".format(job_id), fg="green"))29    start_at = time.perf_counter()30    indexing_cache_key = "app_annotation_batch_import_{}".format(str(job_id))31    # get app info32    app = db.session.query(App).filter(App.id == app_id, App.tenant_id == tenant_id, App.status == "normal").first()33 34    if app:35        try:36            documents = []37            for content in content_list:38                annotation = MessageAnnotation(39                    app_id=app.id, content=content["answer"], question=content["question"], account_id=user_id40                )41                db.session.add(annotation)42                db.session.flush()43 44                document = Document(45                    page_content=content["question"],46                    metadata={"annotation_id": annotation.id, "app_id": app_id, "doc_id": annotation.id},47                )48                documents.append(document)49            # if annotation reply is enabled , batch add annotations' index50            app_annotation_setting = (51                db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()52            )53 54            if app_annotation_setting:55                dataset_collection_binding = (56                    DatasetCollectionBindingService.get_dataset_collection_binding_by_id_and_type(57                        app_annotation_setting.collection_binding_id, "annotation"58                    )59                )60                if not dataset_collection_binding:61                    raise NotFound("App annotation setting not found")62                dataset = Dataset(63                    id=app_id,64                    tenant_id=tenant_id,65                    indexing_technique="high_quality",66                    embedding_model_provider=dataset_collection_binding.provider_name,67                    embedding_model=dataset_collection_binding.model_name,68                    collection_binding_id=dataset_collection_binding.id,69                )70 71                vector = Vector(dataset, attributes=["doc_id", "annotation_id", "app_id"])72                vector.create(documents, duplicate_check=True)73 74            db.session.commit()75            redis_client.setex(indexing_cache_key, 600, "completed")76            end_at = time.perf_counter()77            logging.info(78                click.style(79                    "Build index successful for batch import annotation: {} latency: {}".format(80                        job_id, end_at - start_at81                    ),82                    fg="green",83                )84            )85        except Exception as e:86            db.session.rollback()87            redis_client.setex(indexing_cache_key, 600, "error")88            indexing_error_msg_key = "app_annotation_batch_import_error_msg_{}".format(str(job_id))89            redis_client.setex(indexing_error_msg_key, 600, str(e))90            logging.exception("Build index for batch import annotations failed")91