dify/api/services/knowledge/indexing/execution.py

"""Document indexing phases and their application-owned persistence and execution ports."""

import logging
import time
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import Literal, Protocol

from core.credit_usage import CreditUsageCreatedBy
from core.model_context import with_credit_usage_created_by
from core.rag.models.document import Document
from services.knowledge.indexing.errors import DocumentIsDeletedPausedError, DocumentIsPausedError
from services.knowledge.indexing.estimate import StoredSource
from services.knowledge.resource_scope import DocumentRef

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class IndexingDocument:
    source: StoredSource
    processing_rule: Mapping[str, object]
    doc_language: str
    need_summary: bool

    @property
    def ref(self) -> DocumentRef:
        return self.source.document_ref

    @property
    def doc_form(self) -> str:
        return self.source.document_model


class IndexingDocumentStore(Protocol):
    def get_indexing_document(self, ref: DocumentRef) -> IndexingDocument | None: ...

    def mark_splitting(self, ref: DocumentRef) -> None: ...

    def complete_indexing(self, ref: DocumentRef, *, tokens: int, latency: float) -> None: ...

    def fail_indexing(self, ref: DocumentRef, error: str) -> None: ...


class IndexingSegmentStore(Protocol):
    def clear_for_indexing(self, document: IndexingDocument) -> None: ...

    def save_for_indexing(
        self, document: IndexingDocument, chunks: list[Document], token_counts: list[int]
    ) -> None: ...

    def resume_indexing(self, document: IndexingDocument) -> tuple[list[Document], int]: ...


class IndexingBackend(Protocol):
    def describe_error(self, error: Exception) -> str:
        """Translate backend-specific exceptions into the persisted error message."""
        ...

    def check_paused(self, ref: DocumentRef) -> None: ...

    def extract(self, document: IndexingDocument) -> list[Document]: ...

    def transform(self, document: IndexingDocument, texts: list[Document]) -> list[Document]: ...

    def ensure_admission(self, document: IndexingDocument, chunks: list[Document]) -> None: ...

    def count_tokens(self, document: IndexingDocument, chunks: list[Document]) -> list[int]: ...

    def load(self, document: IndexingDocument, chunks: list[Document]) -> None: ...

    def discard(self, document: IndexingDocument, chunks: list[Document]) -> None:
        """Remove index entries that ``load`` may have written for a document deleted meanwhile."""
        ...


class DocumentIndexingService:
    """Coordinate committed phases; no caller session is held across execution."""

    def __init__(
        self,
        *,
        documents: IndexingDocumentStore,
        segments: IndexingSegmentStore,
        backend: IndexingBackend,
        enforce_vector_space_admission: bool = False,
        clock: Callable[[], float] = time.perf_counter,
    ) -> None:
        self._documents = documents
        self._segments = segments
        self._backend = backend
        self._enforce_vector_space_admission = enforce_vector_space_admission
        self._clock = clock

    @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING)
    def run(self, document_refs: Sequence[DocumentRef]) -> None:
        for ref in document_refs:
            self._run(ref, start="parsing")

    @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING)
    def run_in_splitting_status(self, ref: DocumentRef) -> None:
        self._run(ref, start="splitting")

    @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING)
    def run_in_indexing_status(self, ref: DocumentRef) -> None:
        self._run(ref, start="indexing")

    def _run(self, ref: DocumentRef, *, start: Literal["parsing", "splitting", "indexing"]) -> None:
        # None: no index write has started, so a deletion leaves nothing behind.
        # Set: load may have written these chunks; document cleanup snapshots segment
        # node ids before deleting them and cannot see entries written after its snapshot,
        # so whichever way this run stops, the entries are discarded if the document is gone.
        loading: tuple[IndexingDocument, list[Document]] | None = None
        try:
            document = self._documents.get_indexing_document(ref)
            if document is None:
                logger.info("Document not found, skipping indexing: %s", ref)
                return
            self._backend.check_paused(ref)
            if start == "indexing":
                chunks, total_tokens = self._segments.resume_indexing(document)
            else:
                if start == "splitting":
                    self._segments.clear_for_indexing(document)
                texts = self._backend.extract(document)
                self._documents.mark_splitting(ref)
                chunks = self._backend.transform(document, texts)
                # Recovery retains its existing admission policy; only normal
                # indexing callers that explicitly opt in reserve vector space.
                if start == "parsing" and self._enforce_vector_space_admission:
                    self._backend.ensure_admission(document, chunks)
                token_counts = self._backend.count_tokens(document, chunks)
                total_tokens = sum(token_counts)
                self._segments.save_for_indexing(document, chunks, token_counts)

            started = self._clock()
            self._backend.check_paused(ref)
            loading = (document, chunks)
            self._backend.load(document, chunks)
            self._backend.check_paused(ref)
            self._documents.complete_indexing(ref, tokens=total_tokens, latency=self._clock() - started)
        except DocumentIsPausedError:
            # The pause flag is checked before document existence, so a paused document
            # deleted mid-write surfaces here rather than as a deletion.
            self._discard_if_deleted(ref, loading)
            raise
        except DocumentIsDeletedPausedError:
            logger.info("Document deleted during indexing: %s", ref)
            self._discard(ref, loading)
        except Exception as error:
            logger.exception("Document indexing failed: %s", ref)
            self._documents.fail_indexing(ref, self._backend.describe_error(error))
            self._discard_if_deleted(ref, loading)

    def _discard(self, ref: DocumentRef, loading: tuple[IndexingDocument, list[Document]] | None) -> None:
        """Best-effort removal of entries ``load`` may have written; failures never mask the outcome."""
        if loading is None:
            return
        try:
            self._backend.discard(*loading)
        except Exception:
            logger.exception("Failed to discard index entries of document deleted during indexing: %s", ref)

    def _discard_if_deleted(self, ref: DocumentRef, loading: tuple[IndexingDocument, list[Document]] | None) -> None:
        """Discard written entries when indexing stopped for another reason and the document is gone."""
        if loading is None:
            return
        try:
            self._documents.get_indexing_document(ref)
        except DocumentIsDeletedPausedError:
            logger.info("Document deleted during indexing: %s", ref)
            self._discard(ref, loading)
        except DocumentIsPausedError:
            return
        except Exception:
            logger.warning(
                "Could not check whether document still exists after indexing stopped: %s", ref, exc_info=True
            )