Underground-Digital/Workflow-Engine

1import datetime2import logging3import time4 5import click6from celery import shared_task7from werkzeug.exceptions import NotFound8 9from core.rag.datasource.vdb.vector_factory import Vector10from core.rag.models.document import Document11from extensions.ext_database import db12from extensions.ext_redis import redis_client13from models.dataset import Dataset14from models.model import App, AppAnnotationSetting, MessageAnnotation15from services.dataset_service import DatasetCollectionBindingService16 17 18@shared_task(queue="dataset")19def enable_annotation_reply_task(20    job_id: str,21    app_id: str,22    user_id: str,23    tenant_id: str,24    score_threshold: float,25    embedding_provider_name: str,26    embedding_model_name: str,27):28    """29    Async enable annotation reply task30    """31    logging.info(click.style("Start add app annotation to index: {}".format(app_id), fg="green"))32    start_at = time.perf_counter()33    # get app info34    app = db.session.query(App).filter(App.id == app_id, App.tenant_id == tenant_id, App.status == "normal").first()35 36    if not app:37        raise NotFound("App not found")38 39    annotations = db.session.query(MessageAnnotation).filter(MessageAnnotation.app_id == app_id).all()40    enable_app_annotation_key = "enable_app_annotation_{}".format(str(app_id))41    enable_app_annotation_job_key = "enable_app_annotation_job_{}".format(str(job_id))42 43    try:44        documents = []45        dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(46            embedding_provider_name, embedding_model_name, "annotation"47        )48        annotation_setting = (49            db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()50        )51        if annotation_setting:52            annotation_setting.score_threshold = score_threshold53            annotation_setting.collection_binding_id = dataset_collection_binding.id54            annotation_setting.updated_user_id = user_id55            annotation_setting.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)56            db.session.add(annotation_setting)57        else:58            new_app_annotation_setting = AppAnnotationSetting(59                app_id=app_id,60                score_threshold=score_threshold,61                collection_binding_id=dataset_collection_binding.id,62                created_user_id=user_id,63                updated_user_id=user_id,64            )65            db.session.add(new_app_annotation_setting)66 67        dataset = Dataset(68            id=app_id,69            tenant_id=tenant_id,70            indexing_technique="high_quality",71            embedding_model_provider=embedding_provider_name,72            embedding_model=embedding_model_name,73            collection_binding_id=dataset_collection_binding.id,74        )75        if annotations:76            for annotation in annotations:77                document = Document(78                    page_content=annotation.question,79                    metadata={"annotation_id": annotation.id, "app_id": app_id, "doc_id": annotation.id},80                )81                documents.append(document)82 83            vector = Vector(dataset, attributes=["doc_id", "annotation_id", "app_id"])84            try:85                vector.delete_by_metadata_field("app_id", app_id)86            except Exception as e:87                logging.info(click.style("Delete annotation index error: {}".format(str(e)), fg="red"))88            vector.create(documents)89        db.session.commit()90        redis_client.setex(enable_app_annotation_job_key, 600, "completed")91        end_at = time.perf_counter()92        logging.info(93            click.style("App annotations added to index: {} latency: {}".format(app_id, end_at - start_at), fg="green")94        )95    except Exception as e:96        logging.exception("Annotation batch created index failed:{}".format(str(e)))97        redis_client.setex(enable_app_annotation_job_key, 600, "error")98        enable_app_annotation_error_key = "enable_app_annotation_error_{}".format(str(job_id))99        redis_client.setex(enable_app_annotation_error_key, 600, str(e))100        db.session.rollback()101    finally:102        redis_client.delete(enable_app_annotation_key)103