feat: Batch cancel tasks

1Panel-dev · Dec 24, 2024 · a1d872a · a1d872a
1 parent 0cd9c8f
commit a1d872a
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 1 deletion.
diff --git a/apps/dataset/serializers/document_serializers.py b/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)

diff --git a/apps/dataset/swagger_api/document_api.py b/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)
+                }
+            )
diff --git a/apps/dataset/urls.py b/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(

diff --git a/apps/dataset/views/document.py b/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]