dify/api/tasks/delete_segment_from_index_task.py

import logging
import time

import click
from celery import shared_task
from sqlalchemy import delete, select

from core.db.session_factory import session_factory
from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
from extensions.ext_storage import storage
from models.dataset import Dataset, Document, SegmentAttachmentBinding
from models.model import UploadFile

logger = logging.getLogger(__name__)


@shared_task(queue="dataset")
def delete_segment_from_index_task(
    index_node_ids: list, dataset_id: str, document_id: str, segment_ids: list, child_node_ids: list | None = None
):
    """
    Async Remove segment from index
    :param index_node_ids:
    :param dataset_id:
    :param document_id:

    Usage: delete_segment_from_index_task.delay(index_node_ids, dataset_id, document_id)
    """
    logger.info(click.style("Start delete segment from index", fg="green"))
    start_at = time.perf_counter()
    with session_factory.create_session() as session:
        try:
            dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id).limit(1))
            if not dataset:
                logging.warning("Dataset %s not found, skipping index cleanup", dataset_id)
                return

            dataset_document = session.scalar(select(Document).where(Document.id == document_id).limit(1))
            if not dataset_document:
                return

            if (
                not dataset_document.enabled
                or dataset_document.archived
                or dataset_document.indexing_status != "completed"
            ):
                logging.info("Document not in valid state for index operations, skipping")
                return
            doc_form = dataset_document.doc_form

            # Proceed with index cleanup using the index_node_ids directly
            # For actual deletion, we should delete summaries (not just disable them)
            index_processor = IndexProcessorFactory(doc_form).init_index_processor()
            index_processor.clean(
                dataset,
                index_node_ids,
                with_keywords=True,
                delete_child_chunks=True,
                precomputed_child_node_ids=child_node_ids,
                delete_summaries=True,  # Actually delete summaries when segment is deleted,
                session=session,
            )
            session.commit()
            if dataset.is_multimodal:
                # delete segment attachment binding
                segment_attachment_bindings = session.scalars(
                    select(SegmentAttachmentBinding).where(
                        SegmentAttachmentBinding.tenant_id == dataset.tenant_id,
                        SegmentAttachmentBinding.dataset_id == dataset.id,
                        SegmentAttachmentBinding.document_id == document_id,
                        SegmentAttachmentBinding.segment_id.in_(segment_ids),
                    )
                ).all()
                if segment_attachment_bindings:
                    segment_attachment_bind_ids = [i.id for i in segment_attachment_bindings]
                    attachment_ids = list(
                        dict.fromkeys(binding.attachment_id for binding in segment_attachment_bindings)
                    )

                    for i in range(0, len(segment_attachment_bind_ids), 1000):
                        segment_attachment_bind_delete_stmt = delete(SegmentAttachmentBinding).where(
                            SegmentAttachmentBinding.tenant_id == dataset.tenant_id,
                            SegmentAttachmentBinding.dataset_id == dataset.id,
                            SegmentAttachmentBinding.document_id == document_id,
                            SegmentAttachmentBinding.id.in_(segment_attachment_bind_ids[i : i + 1000]),
                        )
                        session.execute(segment_attachment_bind_delete_stmt)

                    session.flush()
                    remaining_attachment_ids = set(
                        session.scalars(
                            select(SegmentAttachmentBinding.attachment_id).where(
                                SegmentAttachmentBinding.attachment_id.in_(attachment_ids)
                            )
                        ).all()
                    )
                    orphan_attachment_ids = [
                        attachment_id
                        for attachment_id in attachment_ids
                        if attachment_id not in remaining_attachment_ids
                    ]

                    if orphan_attachment_ids:
                        attachment_storage_keys = list(
                            dict.fromkeys(
                                session.scalars(
                                    select(UploadFile.key).where(
                                        UploadFile.tenant_id == dataset.tenant_id,
                                        UploadFile.id.in_(orphan_attachment_ids),
                                    )
                                ).all()
                            )
                        )
                        index_processor.clean(
                            session=session, dataset=dataset, node_ids=orphan_attachment_ids, with_keywords=False
                        )
                        session.execute(
                            delete(UploadFile).where(
                                UploadFile.tenant_id == dataset.tenant_id,
                                UploadFile.id.in_(orphan_attachment_ids),
                            )
                        )
                    else:
                        attachment_storage_keys = []

                    session.commit()

                    for storage_key in attachment_storage_keys:
                        try:
                            storage.delete(storage_key)
                        except Exception:
                            logger.exception("Failed to delete segment attachment from storage, key: %s", storage_key)

            end_at = time.perf_counter()
            logger.info(click.style(f"Segment deleted from index latency: {end_at - start_at}", fg="green"))
        except Exception:
            session.rollback()
            logger.exception("delete segment from index failed")