Underground-Digital/Workflow-Engine

1import logging2import time3 4import click5from celery import shared_task6 7from core.rag.datasource.vdb.vector_factory import Vector8from core.rag.models.document import Document9from models.dataset import Dataset10from services.dataset_service import DatasetCollectionBindingService11 12 13@shared_task(queue="dataset")14def add_annotation_to_index_task(15    annotation_id: str, question: str, tenant_id: str, app_id: str, collection_binding_id: str16):17    """18    Add annotation to index.19    :param annotation_id: annotation id20    :param question: question21    :param tenant_id: tenant id22    :param app_id: app id23    :param collection_binding_id: embedding binding id24 25    Usage: clean_dataset_task.delay(dataset_id, tenant_id, indexing_technique, index_struct)26    """27    logging.info(click.style("Start build index for annotation: {}".format(annotation_id), fg="green"))28    start_at = time.perf_counter()29 30    try:31        dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding_by_id_and_type(32            collection_binding_id, "annotation"33        )34        dataset = Dataset(35            id=app_id,36            tenant_id=tenant_id,37            indexing_technique="high_quality",38            embedding_model_provider=dataset_collection_binding.provider_name,39            embedding_model=dataset_collection_binding.model_name,40            collection_binding_id=dataset_collection_binding.id,41        )42 43        document = Document(44            page_content=question, metadata={"annotation_id": annotation_id, "app_id": app_id, "doc_id": annotation_id}45        )46        vector = Vector(dataset, attributes=["doc_id", "annotation_id", "app_id"])47        vector.create([document], duplicate_check=True)48 49        end_at = time.perf_counter()50        logging.info(51            click.style(52                "Build index successful for annotation: {} latency: {}".format(annotation_id, end_at - start_at),53                fg="green",54            )55        )56    except Exception:57        logging.exception("Build index for annotation failed")58