paperless-ngx/src/documents/caching.py

from __future__ import annotations

import hashlib
import logging
import uuid
from binascii import hexlify
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import Final

from django.core.cache import cache

from documents.models import Document

if TYPE_CHECKING:
    from documents.classifier import DocumentClassifier

logger = logging.getLogger("paperless.caching")


@dataclass(frozen=True)
class MetadataCacheData:
    original_checksum: str
    original_metadata: list
    archive_checksum: str | None
    archive_metadata: list | None


@dataclass(frozen=True)
class SuggestionCacheData:
    classifier_version: int
    classifier_hash: str
    suggestions: dict


CLASSIFIER_VERSION_KEY: Final[str] = "classifier_version"
CLASSIFIER_HASH_KEY: Final[str] = "classifier_hash"
CLASSIFIER_MODIFIED_KEY: Final[str] = "classifier_modified"
# Marker distinguishing LLM suggestions from classifier-generated ones (whose
# FORMAT_VERSION lives in a much lower range - see DocumentClassifier). Bump
# this whenever cached suggestions must not be reused, including changes to
# their shape or interpretation, so a previous release's result cannot leak
# incompatible or obsolete behavior into the new one:
#   1000 - initial LLM suggestions cache (flat lists of resolved object ids
#          per taxonomy field)
#   1001 - suggestions reshaped to {"existing_ids": [...], "new_names":
#          [...]} per taxonomy field (#13676)
#   1002 - names are always generated and optional candidate mappings are
#          validated separately, so candidate-anchored 1001 results are stale
LLM_CACHE_CLASSIFIER_VERSION: Final[int] = 1002

CACHE_1_MINUTE: Final[int] = 60
CACHE_5_MINUTES: Final[int] = 5 * CACHE_1_MINUTE
CACHE_50_MINUTES: Final[int] = 50 * CACHE_1_MINUTE
# Deliberately longer than any entry it names
LLM_CACHE_GENERATION_TIMEOUT: Final[int] = 2 * CACHE_50_MINUTES


def get_suggestion_cache_key(document_id: int) -> str:
    """
    Returns the basic key for a document's suggestions
    """
    return f"doc_{document_id}_suggest"


def get_suggestion_cache(document_id: int) -> SuggestionCacheData | None:
    """
    If possible, return the cached suggestions for the given document ID.
    The classifier needs to be matching in format and hash and the suggestions need to
    have been cached once.
    """
    from documents.classifier import DocumentClassifier

    doc_key = get_suggestion_cache_key(document_id)
    cache_hits = cache.get_many([CLASSIFIER_VERSION_KEY, CLASSIFIER_HASH_KEY, doc_key])
    # The document suggestions are in the cache
    if doc_key in cache_hits:
        doc_suggestions: SuggestionCacheData = cache_hits[doc_key]
        # The classifier format is the same
        # The classifier hash is the same
        # Then the suggestions can be used
        if (
            CLASSIFIER_VERSION_KEY in cache_hits
            and cache_hits[CLASSIFIER_VERSION_KEY] == DocumentClassifier.FORMAT_VERSION
            and cache_hits[CLASSIFIER_VERSION_KEY] == doc_suggestions.classifier_version
        ) and (
            CLASSIFIER_HASH_KEY in cache_hits
            and cache_hits[CLASSIFIER_HASH_KEY] == doc_suggestions.classifier_hash
        ):
            return doc_suggestions
        else:  # pragma: no cover
            # Remove the key because something didn't match
            cache.delete(doc_key)
    return None


def set_suggestions_cache(
    document_id: int,
    suggestions: dict,
    classifier: DocumentClassifier | None,
    *,
    timeout=CACHE_50_MINUTES,
) -> None:
    """
    Caches the given suggestions, which were generated by the given classifier.  If there is no classifier,
    this function is a no-op (there won't be suggestions then anyway)
    """
    if classifier is not None:
        doc_key = get_suggestion_cache_key(document_id)
        cache.set(
            doc_key,
            SuggestionCacheData(
                classifier.FORMAT_VERSION,
                hexlify(classifier.last_auto_type_hash).decode(),
                suggestions,
            ),
            timeout,
        )


def refresh_suggestions_cache(
    document_id: int,
    *,
    timeout: int = CACHE_50_MINUTES,
) -> None:
    """
    Refreshes the expiration of the suggestions for the given document ID
    to the given timeout
    """
    doc_key = get_suggestion_cache_key(document_id)
    cache.touch(doc_key, timeout)


def invalidate_suggestions_cache(document_id: int) -> None:
    """Invalidate classifier-generated suggestions for a document."""
    cache.delete(get_suggestion_cache_key(document_id))


def _llm_generation_key(document_id: int) -> str:
    return f"{get_suggestion_cache_key(document_id)}_llm_generation"


def _llm_variant_key(document_id: int, backend: str) -> str:
    """Cache key for one LLM configuration and permission scope.

    ``backend`` identifies the variant - model, endpoint, output language and
    requesting user.

    Generating the token on first use lets invalidate_llm_suggestions_cache()
    be no-op for documents that never had AI suggestions.
    """
    generation_key = _llm_generation_key(document_id)
    generation = cache.get_or_set(
        generation_key,
        lambda: uuid.uuid4().hex,
        timeout=LLM_CACHE_GENERATION_TIMEOUT,
    )
    cache.touch(generation_key, LLM_CACHE_GENERATION_TIMEOUT)
    backend_hash = hashlib.sha256(backend.encode()).hexdigest()[:16]
    return f"{get_suggestion_cache_key(document_id)}_llm_{generation}_{backend_hash}"


def get_llm_suggestion_cache(
    document_id: int,
    backend: str,
) -> SuggestionCacheData | None:
    data: SuggestionCacheData = cache.get(_llm_variant_key(document_id, backend))

    if (
        data
        and data.classifier_version == LLM_CACHE_CLASSIFIER_VERSION
        and data.classifier_hash == backend
    ):
        return data

    return None


def set_llm_suggestions_cache(
    document_id: int,
    suggestions: dict,
    *,
    backend: str,
    timeout: int = CACHE_50_MINUTES,
) -> None:
    """
    Cache LLM-generated suggestions using a backend-specific identifier
    (e.g. 'openai-like:gpt-4').
    """
    cache.set(
        _llm_variant_key(document_id, backend),
        SuggestionCacheData(
            classifier_version=LLM_CACHE_CLASSIFIER_VERSION,
            classifier_hash=backend,
            suggestions=suggestions,
        ),
        timeout,
    )


def refresh_llm_suggestions_cache(
    document_id: int,
    backend: str,
    *,
    timeout: int = CACHE_50_MINUTES,
) -> None:
    """
    Refreshes the expiration of one cached LLM suggestion variant.
    """
    cache.touch(_llm_variant_key(document_id, backend), timeout)


def invalidate_llm_suggestions_cache(
    document_id: int,
) -> None:
    """
    Invalidate every LLM suggestion variant for a document.
    """
    generation_key = _llm_generation_key(document_id)
    if cache.get(generation_key) is not None:
        cache.set(
            generation_key,
            uuid.uuid4().hex,
            timeout=LLM_CACHE_GENERATION_TIMEOUT,
        )


def get_metadata_cache_key(document_id: int) -> str:
    """
    Returns the basic key for a document's metadata
    """
    return f"doc_{document_id}_metadata"


def get_metadata_cache(document_id: int) -> MetadataCacheData | None:
    """
    Returns the cached document metadata for the given document ID, as long as the metadata
    was cached once and the checksums have not changed
    """
    doc_key = get_metadata_cache_key(document_id)
    doc_metadata: MetadataCacheData | None = cache.get(doc_key)
    # The metadata exists in the cache
    if doc_metadata is not None:
        try:
            doc = Document.objects.only(
                "pk",
                "checksum",
                "archive_checksum",
                "archive_filename",
            ).get(pk=document_id)
            # The original checksums match
            # If it has one, the archive checksums match
            # Then, we can use the metadata
            if (
                doc_metadata.original_checksum == doc.checksum
                and doc.has_archive_version
                and doc_metadata.archive_checksum is not None
                and doc_metadata.archive_checksum == doc.archive_checksum
            ):
                # Refresh cache
                cache.touch(doc_key, CACHE_50_MINUTES)
                return doc_metadata
            else:  # pragma: no cover
                # Something didn't match, delete the key
                cache.delete(doc_key)
        except Document.DoesNotExist:  # pragma: no cover
            # Basically impossible, but the key existed, but the Document didn't
            cache.delete(doc_key)
    return None


def set_metadata_cache(
    document: Document,
    original_metadata: list,
    archive_metadata: list | None,
    *,
    timeout=CACHE_50_MINUTES,
) -> None:
    """
    Sets the metadata into cache for the given Document
    """
    doc_key = get_metadata_cache_key(document.pk)
    cache.set(
        doc_key,
        MetadataCacheData(
            document.checksum,
            original_metadata,
            document.archive_checksum,
            archive_metadata,
        ),
        timeout,
    )


def refresh_metadata_cache(
    document_id: int,
    *,
    timeout: int = CACHE_50_MINUTES,
) -> None:
    """
    Refreshes the expiration of the metadata for the given document ID
    to the given timeout
    """
    doc_key = get_metadata_cache_key(document_id)
    cache.touch(doc_key, timeout)


def get_thumbnail_modified_key(document_id: int) -> str:
    """
    Builds the key to store a thumbnail's timestamp
    """
    return f"doc_{document_id}_thumbnail_modified"


def clear_document_caches(document_id: int) -> None:
    """
    Removes all cached items for the given document
    """
    cache.delete_many(
        [
            get_suggestion_cache_key(document_id),
            get_metadata_cache_key(document_id),
            get_thumbnail_modified_key(document_id),
        ],
    )
    invalidate_llm_suggestions_cache(document_id)