dify/api/services/knowledge/indexing/adapters/cleanup.py

"""Delete document indexes between bounded database transactions."""

from collections.abc import Callable, Mapping, Sequence

from sqlalchemy import delete, select
from sqlalchemy.orm import Session

from configs import dify_config
from core.rag.datasource.keyword.jieba.jieba import Jieba
from core.rag.datasource.vdb.vector_factory import Vector
from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from extensions.ext_redis import redis_client
from models.dataset import ChildChunk, Dataset, DocumentSegment, DocumentSegmentSummary
from repositories.knowledge.dataset_read_repository import get_dataset_keyword_table
from repositories.knowledge.keyword_table_repository import persist_keyword_table


def delete_index_entries(
    dataset: Dataset,
    *,
    vector_node_ids: Sequence[str],
    vector_type: str | None,
    keyword_node_ids: Sequence[str],
    new_session: Callable[[], Session],
) -> None:
    """Delete index entries by node id; the single owner of how each index type is removed.

    High-quality datasets delete ``vector_node_ids`` from the vector store; other datasets
    delete ``keyword_node_ids`` from the keyword table under the dataset keyword lock.
    """
    if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
        if vector_node_ids:
            Vector(dataset, session=None, vector_type=vector_type).delete_by_ids(list(vector_node_ids))
        return
    if not keyword_node_ids:
        return

    # The keyword adapter owns its lock and storage I/O; callbacks only do SQL.
    def read() -> tuple[str, str | None]:
        with new_session() as session:
            row = get_dataset_keyword_table(dataset, session=session)
            return (row.data_source_type, row.keyword_table) if row else (dify_config.KEYWORD_DATA_SOURCE_TYPE, None)

    def write(storage_type: str, data: str, keywords: Mapping[str, Sequence[str]]) -> None:
        with new_session() as session, session.begin():
            persist_keyword_table(
                session,
                tenant_id=dataset.tenant_id,
                dataset_id=dataset.id,
                storage_type=storage_type,
                data=data,
                keywords=keywords,
            )

    Jieba(dataset).update_texts(
        [],
        read=read,
        write=write,
        lock=redis_client.lock(f"keyword_indexing_lock_{dataset.id}", timeout=600),
        deleted_ids=keyword_node_ids,
    )


def clean_document_indexes(
    *, dataset_id: str, document_ids: Sequence[str], doc_form: str | None, new_session: Callable[[], Session]
) -> str | None:
    """Clean indexes even when the owning document rows have already been deleted.

    Snapshot the dataset-owned rows, close the read session, remove external
    indexes, then delete the selected summary and child rows in one transaction.
    A failed vector deletion leaves those rows available for a retry.
    Empty document selections never mean deleting the entire dataset index.
    """
    if not document_ids:
        return None
    with new_session() as session:
        dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id))
        if dataset is None:
            return None
        tenant_id = dataset.tenant_id
        segments = session.scalars(
            select(DocumentSegment).where(
                DocumentSegment.tenant_id == tenant_id,
                DocumentSegment.dataset_id == dataset_id,
                DocumentSegment.document_id.in_(document_ids),
            )
        ).all()
        summaries = session.scalars(
            select(DocumentSegmentSummary).where(
                DocumentSegmentSummary.dataset_id == dataset_id,
                DocumentSegmentSummary.document_id.in_(document_ids),
            )
        ).all()
        children = session.scalars(
            select(ChildChunk).where(
                ChildChunk.tenant_id == tenant_id,
                ChildChunk.dataset_id == dataset_id,
                ChildChunk.document_id.in_(document_ids),
            )
        ).all()
        summary_ids = [summary.id for summary in summaries]
        child_ids = [child.id for child in children]
        summary_node_ids = [summary.summary_index_node_id for summary in summaries if summary.summary_index_node_id]
        body_node_ids = (
            [child.index_node_id for child in children if child.index_node_id]
            if doc_form == IndexStructureType.PARENT_CHILD_INDEX
            else [segment.index_node_id for segment in segments if segment.index_node_id]
        )
        high_quality = dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY
        node_ids = list(dict.fromkeys([*summary_node_ids, *body_node_ids]))
        vector_type = Vector.resolve_vector_type(dataset, session=session) if high_quality and node_ids else None
        session.expunge(dataset)

    if not node_ids and not summary_ids and not child_ids:
        return None

    delete_index_entries(
        dataset,
        vector_node_ids=node_ids,
        vector_type=vector_type,
        keyword_node_ids=body_node_ids,
        new_session=new_session,
    )

    with new_session() as session, session.begin():
        if summary_ids:
            session.execute(
                delete(DocumentSegmentSummary).where(
                    DocumentSegmentSummary.dataset_id == dataset_id,
                    DocumentSegmentSummary.document_id.in_(document_ids),
                    DocumentSegmentSummary.id.in_(summary_ids),
                )
            )
        if child_ids:
            session.execute(
                delete(ChildChunk).where(
                    ChildChunk.tenant_id == tenant_id,
                    ChildChunk.dataset_id == dataset_id,
                    ChildChunk.document_id.in_(document_ids),
                    ChildChunk.id.in_(child_ids),
                )
            )
    return tenant_id