From 9e76cd97de0f84b9afcb5340ca91747e1258597d Mon Sep 17 00:00:00 2001
From: shaohuzhang1 <shaohu.zhang@fit2cloud.com>
Date: Thu, 18 Jul 2024 16:36:34 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E5=90=91=E9=87=8F?=
 =?UTF-8?q?=E6=A8=A1=E5=9E=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/common/event/listener_manage.py            | 4 ++++
 apps/dataset/serializers/dataset_serializers.py | 1 +
 2 files changed, 5 insertions(+)

diff --git a/apps/common/event/listener_manage.py b/apps/common/event/listener_manage.py
index fc94fd838..1e874062c 100644
--- a/apps/common/event/listener_manage.py
+++ b/apps/common/event/listener_manage.py
@@ -159,6 +159,8 @@ class ListenerManagement:
         @param embedding_model 向量模型
         :return: None
         """
+        if not try_lock('embedding' + document_id):
+            return
         max_kb.info(f"开始--->向量化文档:{document_id}")
         QuerySet(Document).filter(id=document_id).update(**{'status': Status.embedding})
         QuerySet(Paragraph).filter(document_id=document_id).update(**{'status': Status.embedding})
@@ -184,6 +186,7 @@ class ListenerManagement:
                 **{'status': status, 'update_time': datetime.datetime.now()})
             QuerySet(Paragraph).filter(document_id=document_id).update(**{'status': status})
             max_kb.info(f"结束--->向量化文档:{document_id}")
+            un_lock('embedding' + document_id)
 
     @staticmethod
     @embedding_poxy
@@ -196,6 +199,7 @@ class ListenerManagement:
         """
         max_kb.info(f"开始--->向量化数据集:{dataset_id}")
         try:
+            ListenerManagement.delete_embedding_by_dataset(dataset_id)
             document_list = QuerySet(Document).filter(dataset_id=dataset_id)
             max_kb.info(f"数据集文档:{[d.name for d in document_list]}")
             for document in document_list:
diff --git a/apps/dataset/serializers/dataset_serializers.py b/apps/dataset/serializers/dataset_serializers.py
index 521aa661a..946398482 100644
--- a/apps/dataset/serializers/dataset_serializers.py
+++ b/apps/dataset/serializers/dataset_serializers.py
@@ -785,6 +785,7 @@ class DataSetSerializers(serializers.ModelSerializer):
                             QuerySet(ApplicationDatasetMapping).filter(
                                 dataset_id=self.data.get('id'))]))}
 
+        @transaction.atomic
         def edit(self, dataset: Dict, user_id: str):
             """
             修改知识库