Annotation management (#1767)

Co-authored-by: jyong <jyong@dify.ai>
2023-12-18 13:10:05 +08:00
parent a9b942981d
commit a71f2863ac
41 changed files with 1871 additions and 67 deletions
--- a/api/tasks/annotation/disable_annotation_reply_task.py
+++ b/api/tasks/annotation/disable_annotation_reply_task.py
@@ -0,0 +1,74 @@
+import datetime
+import logging
+import time
+
+import click
+from celery import shared_task
+from werkzeug.exceptions import NotFound
+
+from core.index.index import IndexBuilder
+from extensions.ext_database import db
+from extensions.ext_redis import redis_client
+from models.dataset import Dataset
+from models.model import MessageAnnotation, App, AppAnnotationSetting
+
+
+@shared_task(queue='dataset')
+def disable_annotation_reply_task(job_id: str, app_id: str, tenant_id: str):
+    """
+    Async enable annotation reply task
+    """
+    logging.info(click.style('Start delete app annotations index: {}'.format(app_id), fg='green'))
+    start_at = time.perf_counter()
+    # get app info
+    app = db.session.query(App).filter(
+        App.id == app_id,
+        App.tenant_id == tenant_id,
+        App.status == 'normal'
+    ).first()
+
+    if not app:
+        raise NotFound("App not found")
+
+    app_annotation_setting = db.session.query(AppAnnotationSetting).filter(
+        AppAnnotationSetting.app_id == app_id
+    ).first()
+
+    if not app_annotation_setting:
+        raise NotFound("App annotation setting not found")
+
+    disable_app_annotation_key = 'disable_app_annotation_{}'.format(str(app_id))
+    disable_app_annotation_job_key = 'disable_app_annotation_job_{}'.format(str(job_id))
+
+    try:
+
+        dataset = Dataset(
+            id=app_id,
+            tenant_id=tenant_id,
+            indexing_technique='high_quality',
+            collection_binding_id=app_annotation_setting.collection_binding_id
+        )
+
+        vector_index = IndexBuilder.get_default_high_quality_index(dataset)
+        if vector_index:
+            try:
+                vector_index.delete_by_metadata_field('app_id', app_id)
+            except Exception:
+                logging.exception("Delete doc index failed when dataset deleted.")
+        redis_client.setex(disable_app_annotation_job_key, 600, 'completed')
+
+        # delete annotation setting
+        db.session.delete(app_annotation_setting)
+        db.session.commit()
+
+        end_at = time.perf_counter()
+        logging.info(
+            click.style('App annotations index deleted : {} latency: {}'.format(app_id, end_at - start_at),
+                        fg='green'))
+    except Exception as e:
+        logging.exception("Annotation batch deleted index failed:{}".format(str(e)))
+        redis_client.setex(disable_app_annotation_job_key, 600, 'error')
+        disable_app_annotation_error_key = 'disable_app_annotation_error_{}'.format(str(job_id))
+        redis_client.setex(disable_app_annotation_error_key, 600, str(e))
+    finally:
+        redis_client.delete(disable_app_annotation_key)