disable all chunks status when disable document (#12149)

api/tasks/add_document_to_index_task.py

@@ -38,7 +38,11 @@ def add_document_to_index_task(dataset_document_id: str):
     try:
         segments = (
             db.session.query(DocumentSegment)
-            .filter(DocumentSegment.document_id == dataset_document.id, DocumentSegment.enabled == True)
+            .filter(
+                DocumentSegment.document_id == dataset_document.id,
+                DocumentSegment.enabled == False,
+                DocumentSegment.status == "completed",
+            )
             .order_by(DocumentSegment.position.asc())
             .all()
         )
@@ -85,6 +89,16 @@ def add_document_to_index_task(dataset_document_id: str):
         db.session.query(DatasetAutoDisableLog).filter(
             DatasetAutoDisableLog.document_id == dataset_document.id
         ).delete()
+
+        # update segment to enable
+        db.session.query(DocumentSegment).filter(DocumentSegment.document_id == dataset_document.id).update(
+            {
+                DocumentSegment.enabled: True,
+                DocumentSegment.disabled_at: None,
+                DocumentSegment.disabled_by: None,
+                DocumentSegment.updated_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
+            }
+        )
         db.session.commit()
 
         end_at = time.perf_counter()

api/tasks/remove_document_from_index_task.py

@@ -1,3 +1,4 @@
+import datetime
 import logging
 import time
 
@@ -46,6 +47,16 @@ def remove_document_from_index_task(document_id: str):
                 index_processor.clean(dataset, index_node_ids, with_keywords=True, delete_child_chunks=False)
             except Exception:
                 logging.exception(f"clean dataset {dataset.id} from index failed")
+        # update segment to disable
+        db.session.query(DocumentSegment).filter(DocumentSegment.document_id == document.id).update(
+            {
+                DocumentSegment.enabled: False,
+                DocumentSegment.disabled_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
+                DocumentSegment.disabled_by: document.disabled_by,
+                DocumentSegment.updated_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
+            }
+        )
+        db.session.commit()
 
         end_at = time.perf_counter()
         logging.info(