feat: Batch cancel tasks (#1896)

Author: shaohuzhang1

apps/dataset/serializers/document_serializers.py

@@ -77,8 +77,23 @@ def get_buffer(self, file):
         return self.buffer
 
 
+class BatchCancelInstanceSerializer(serializers.Serializer):
+    id_list = serializers.ListField(required=True, child=serializers.UUIDField(required=True),
+                                    error_messages=ErrMessage.char("id列表"))
+    type = serializers.IntegerField(required=True, error_messages=ErrMessage.integer(
+        "任务类型"))
+
+    def is_valid(self, *, raise_exception=False):
+        super().is_valid(raise_exception=True)
+        _type = self.data.get('type')
+        try:
+            TaskType(_type)
+        except Exception as e:
+            raise AppApiException(500, '任务类型不支持')
+
+
 class CancelInstanceSerializer(serializers.Serializer):
-    type = serializers.IntegerField(required=True, error_messages=ErrMessage.boolean(
+    type = serializers.IntegerField(required=True, error_messages=ErrMessage.integer(
         "任务类型"))
 
     def is_valid(self, *, raise_exception=False):
@@ -1064,6 +1079,28 @@ def batch_delete(self, instance: Dict, with_valid=True):
             delete_embedding_by_document_list(document_id_list)
             return True
 
+        def batch_cancel(self, instance: Dict, with_valid=True):
+            if with_valid:
+                self.is_valid(raise_exception=True)
+                BatchCancelInstanceSerializer(data=instance).is_valid(raise_exception=True)
+            document_id_list = instance.get("id_list")
+            ListenerManagement.update_status(QuerySet(Paragraph).annotate(
+                reversed_status=Reverse('status'),
+                task_type_status=Substr('reversed_status', TaskType(instance.get('type')).value,
+                                        1),
+            ).filter(task_type_status__in=[State.PENDING.value, State.STARTED.value]).filter(
+                document_id__in=document_id_list).values('id'),
+                                             TaskType(instance.get('type')),
+                                             State.REVOKE)
+            ListenerManagement.update_status(QuerySet(Document).annotate(
+                reversed_status=Reverse('status'),
+                task_type_status=Substr('reversed_status', TaskType(instance.get('type')).value,
+                                        1),
+            ).filter(task_type_status__in=[State.PENDING.value, State.STARTED.value]).filter(
+                id__in=document_id_list).values('id'),
+                                             TaskType(instance.get('type')),
+                                             State.REVOKE)
+
         def batch_edit_hit_handling(self, instance: Dict, with_valid=True):
             if with_valid:
                 BatchSerializer(data=instance).is_valid(model=Document, raise_exception=True)

apps/dataset/swagger_api/document_api.py

@@ -37,3 +37,17 @@ def get_request_body_api():
                                            description="1|2|3 1:向量化|2:生成问题|3:同步文档")
                 }
             )
+
+    class BatchCancel(ApiMixin):
+        @staticmethod
+        def get_request_body_api():
+            return openapi.Schema(
+                type=openapi.TYPE_OBJECT,
+                properties={
+                    'id_list': openapi.Schema(type=openapi.TYPE_ARRAY, items=openapi.Schema(type=openapi.TYPE_STRING),
+                                              title="文档id列表",
+                                              description="文档id列表"),
+                    'type': openapi.Schema(type=openapi.TYPE_INTEGER, title="任务类型",
+                                           description="1|2|3 1:向量化|2:生成问题|3:同步文档", default=1)
+                }
+            )

apps/dataset/urls.py

@@ -41,6 +41,8 @@
     path('dataset/<str:dataset_id>/document/<str:document_id>/sync', views.Document.SyncWeb.as_view()),
     path('dataset/<str:dataset_id>/document/<str:document_id>/refresh', views.Document.Refresh.as_view()),
     path('dataset/<str:dataset_id>/document/<str:document_id>/cancel_task', views.Document.CancelTask.as_view()),
+    path('dataset/<str:dataset_id>/document/cancel_task/_batch',
+         views.Document.CancelTask.Batch.as_view()),
     path('dataset/<str:dataset_id>/document/<str:document_id>/paragraph', views.Paragraph.as_view()),
     path('dataset/<str:dataset_id>/document/batch_generate_related', views.Document.BatchGenerateRelated.as_view()),
     path(

apps/dataset/views/document.py

@@ -238,6 +238,24 @@ def put(self, request: Request, dataset_id: str, document_id: str):
                     request.data
                 ))
 
+        class Batch(APIView):
+            authentication_classes = [TokenAuth]
+
+            @action(methods=['PUT'], detail=False)
+            @swagger_auto_schema(operation_summary="批量取消任务",
+                                 operation_id="批量取消任务",
+                                 request_body=DocumentApi.BatchCancel.get_request_body_api(),
+                                 manual_parameters=DocumentSerializers.Create.get_request_params_api(),
+                                 responses=result.get_default_response(),
+                                 tags=["知识库/文档"]
+                                 )
+            @has_permissions(
+                lambda r, k: Permission(group=Group.DATASET, operate=Operate.MANAGE,
+                                        dynamic_tag=k.get('dataset_id')))
+            def put(self, request: Request, dataset_id: str):
+                return result.success(
+                    DocumentSerializers.Batch(data={'dataset_id': dataset_id}).batch_cancel(request.data))
+
     class Refresh(APIView):
         authentication_classes = [TokenAuth]
 

ui/src/api/document.ts

@@ -359,6 +359,14 @@ const cancelTask: (
   )
 }
 
+const batchCancelTask: (
+  dataset_id: string,
+  data: any,
+  loading?: Ref<boolean>
+) => Promise<Result<boolean>> = (dataset_id, data, loading) => {
+  return put(`${prefix}/${dataset_id}/document/cancel_task/_batch`, data, undefined, loading)
+}
+
 export default {
   postSplitDocument,
   getDocument,
@@ -383,5 +391,6 @@ export default {
   batchRefresh,
   batchGenerateRelated,
   cancelTask,
-  exportDocumentZip
+  exportDocumentZip,
+  batchCancelTask
 }

ui/src/styles/element-plus.scss

@@ -6,7 +6,6 @@
   --el-text-color-regular: #1f2329;
   --el-color-info: #8f959e !important;
   --el-disabled-bg-color: #eff0f1;
-  --el-disabled-border-color: #bbbfc4;
   --el-text-color-primary: #1f2329;
 }
 

ui/src/views/document/index.vue

@@ -1,5 +1,5 @@
 <template>
-  <LayoutContainer header="文档">
+  <LayoutContainer header="文档" class="document-main">
     <div class="main-calc-height">
       <div class="p-24">
         <div class="flex-between">
@@ -403,12 +403,25 @@
           </el-table-column>
         </app-table>
       </div>
+
       <ImportDocumentDialog ref="ImportDocumentDialogRef" :title="title" @refresh="refresh" />
       <SyncWebDialog ref="SyncWebDialogRef" @refresh="refresh" />
       <!-- 选择知识库 -->
       <SelectDatasetDialog ref="SelectDatasetDialogRef" @refresh="refreshMigrate" />
       <GenerateRelatedDialog ref="GenerateRelatedDialogRef" @refresh="refresh" />
     </div>
+    <div class="mul-operation w-full flex" v-if="multipleSelection.length !== 0">
+      <el-button :disabled="multipleSelection.length === 0" @click="cancelTaskHandle(1)">
+        取消向量化
+      </el-button>
+      <el-button :disabled="multipleSelection.length === 0" @click="cancelTaskHandle(2)">
+        取消生成
+      </el-button>
+      <el-text type="info" class="secondary ml-24">
+        已选 {{ multipleSelection.length }} 项
+      </el-text>
+      <el-button class="ml-16" type="primary" link @click="clearSelection"> 清空 </el-button>
+    </div>
   </LayoutContainer>
 </template>
 <script setup lang="ts">
@@ -478,6 +491,7 @@ const multipleSelection = ref<any[]>([])
 const title = ref('')
 
 const SelectDatasetDialogRef = ref()
+
 const exportDocument = (document: any) => {
   documentApi.exportDocument(document.name, document.dataset_id, document.id, loading).then(() => {
     MsgSuccess('导出成功')
@@ -490,6 +504,28 @@ const exportDocumentZip = (document: any) => {
       MsgSuccess('导出成功')
     })
 }
+
+function cancelTaskHandle(val: any) {
+  const arr: string[] = []
+  multipleSelection.value.map((v) => {
+    if (v) {
+      arr.push(v.id)
+    }
+  })
+  const obj = {
+    id_list: arr,
+    type: val
+  }
+  documentApi.batchCancelTask(id, obj, loading).then(() => {
+    MsgSuccess('批量取消成功')
+    multipleTableRef.value?.clearSelection()
+  })
+}
+
+function clearSelection() {
+  multipleTableRef.value?.clearSelection()
+}
+
 function openDatasetDialog(row?: any) {
   const arr: string[] = []
   if (row) {
@@ -813,4 +849,20 @@ onBeforeUnmount(() => {
   closeInterval()
 })
 </script>
-<style lang="scss" scoped></style>
+<style lang="scss" scoped>
+.document-main {
+  box-sizing: border-box;
+  .mul-operation {
+    position: fixed;
+    margin-left: var(--sidebar-width);
+    bottom: 0;
+    right: 24px;
+    width: calc(100% - var(--sidebar-width) - 48px);
+    padding: 16px 24px;
+    box-sizing: border-box;
+    background: #ffffff;
+    z-index: 22;
+    box-shadow: 0px -2px 4px 0px rgba(31, 35, 41, 0.08);
+  }
+}
+</style>