paperless-ngx/src/paperless_ai/ai_classifier.py

import json
import logging

from django.conf import settings
from django.contrib.auth.models import User

from documents.models import Document
from documents.permissions import get_objects_for_user_owner_aware
from paperless.config import AIConfig
from paperless_ai.base_model import ClassificationSuggestions
from paperless_ai.base_model import TaxonomyChoiceDict
from paperless_ai.client import AIClient
from paperless_ai.db import db_connection_released
from paperless_ai.indexing import _node_document_ids
from paperless_ai.indexing import retrieve_similar_nodes
from paperless_ai.indexing import truncate_content
from paperless_ai.prompts.context import ClassificationPromptContext
from paperless_ai.prompts.context import LocalizationPromptContext
from paperless_ai.prompts.context import RagContextPromptContext
from paperless_ai.prompts.render import render_prompt
from paperless_ai.taxonomy import AssignedMetadata
from paperless_ai.taxonomy import TaxonomyCandidates
from paperless_ai.taxonomy import build_taxonomy_candidates
from paperless_ai.taxonomy import empty_taxonomy_candidates
from paperless_ai.taxonomy import format_taxonomy_for_prompt
from paperless_ai.taxonomy import get_assigned_metadata

logger = logging.getLogger("paperless_ai.rag_classifier")

# Neighbours retrieved for taxonomy-candidate weighting, decoupled from
# get_taxonomy_context's max_docs (which caps how many of those same
# neighbours get their text spliced into the RAG context block). A wider
# pool of weighted neighbours gives build_taxonomy_candidates() more signal
# for which tags/correspondents/etc. actually cluster around this document,
# while the ranked candidate lists it returns stay capped by
# taxonomy.MAX_TAG_CANDIDATES / MAX_SINGLE_VALUE_CANDIDATES regardless of
# how many neighbours went in - so raising this does not by itself grow the
# prompt.
TAXONOMY_CANDIDATE_TOP_K = 15


def get_language_name(language_code: str) -> str:
    normalized_language_code = language_code.lower()
    for code, name in settings.LANGUAGES:
        if code.lower() == normalized_language_code:
            return str(name)
    return language_code


def build_prompt_without_rag(
    document: Document,
    config: AIConfig,
    candidates: TaxonomyCandidates | None = None,
    assigned: AssignedMetadata | None = None,
) -> str:
    filename = document.filename or ""
    content = truncate_content(
        document.content[:4000] or "",
        chunk_size=config.llm_embedding_chunk_size,
        context_size=config.llm_context_size,
    )

    taxonomy_block = (
        format_taxonomy_for_prompt(candidates, assigned)
        if candidates is not None and assigned is not None
        else ""
    )
    has_candidates = candidates is not None and any(candidates.values())

    return render_prompt(
        ClassificationPromptContext(
            filename=filename,
            content=content,
            taxonomy_block=taxonomy_block,
            has_candidates=has_candidates,
        ),
    )


def build_prompt_with_rag(
    document: Document,
    config: AIConfig,
    candidates: TaxonomyCandidates | None = None,
    assigned: AssignedMetadata | None = None,
    context: str = "",
) -> str:
    base_prompt = build_prompt_without_rag(
        document,
        config,
        candidates=candidates,
        assigned=assigned,
    )
    truncated_context = truncate_content(
        context,
        chunk_size=config.llm_embedding_chunk_size,
        context_size=config.llm_context_size,
    )

    return render_prompt(
        RagContextPromptContext(
            base_prompt=base_prompt,
            context=truncated_context,
        ),
    )


def build_localization_prompt(
    suggestions: ClassificationSuggestions,
    output_language: str,
) -> str:
    """``suggestions`` is the full nested-shape result of parse_ai_response
    (each taxonomy field a ``{"existing_ids": [...], "new_names": [...]}``
    dict) - passed through as-is so the model receives and returns the exact
    DocumentClassifierSchema shape run_llm_query() always parses against.
    Only each field's new_names (never existing_ids, which are plain
    resolved-object IDs, not text) and title get used from the response; see
    get_ai_document_classification's merge step, which always keeps the
    *original* existing_ids regardless of what the model echoes back here.
    """
    language_name = get_language_name(output_language)
    return render_prompt(
        LocalizationPromptContext(
            language_name=language_name,
            suggestions_json=json.dumps(suggestions, ensure_ascii=False),
        ),
    )


def get_taxonomy_context(
    document: Document,
    user: User | None = None,
    max_docs: int = 5,
) -> tuple[TaxonomyCandidates, AssignedMetadata, str]:
    """One retrieval feeds both taxonomy candidates and RAG text context.
    On any retrieval failure, degrades to empty candidates/context rather than
    propagating the exception - a vector-store outage should not block
    classification, only its RAG-assisted enrichment.
    """
    assigned = get_assigned_metadata(document, user)
    try:
        # None means "no restriction" to retrieve_similar_nodes. A superuser
        # (like no user at all) can see every document, so skip materializing
        # every visible pk into a Python list and passing it through as an IN
        # filter: for a large library that is a wasted quadratic scan in the
        # vector store at best, and past ~32,763 documents a hard
        # sqlite3.OperationalError (SQLite's bound-parameter limit) at worst.
        # get_objects_for_user_owner_aware() would return every Document for a
        # superuser anyway (guardian's own with_superuser shortcut), so this
        # changes nothing about which documents are considered -- only how we
        # get there.
        visible_document_ids = (
            None
            if user is None or user.is_superuser
            else list(
                get_objects_for_user_owner_aware(
                    user,
                    "view_document",
                    Document,
                ).values_list("pk", flat=True),
            )
        )
        nodes = retrieve_similar_nodes(
            document,
            top_k=TAXONOMY_CANDIDATE_TOP_K,
            document_ids=visible_document_ids,
        )

        candidates = build_taxonomy_candidates(nodes, user)

        similar_docs = list(
            Document.objects.filter(pk__in=_node_document_ids(nodes))[:max_docs],
        )
        context_blocks = []
        for similar in similar_docs:
            text = similar.content[:1000] or ""
            title = similar.title or similar.filename or "Untitled"
            context_blocks.append(f"TITLE: {title}\n{text}")
    except Exception:
        logger.exception(
            "Failed to retrieve RAG neighbours for document %s; continuing "
            "without taxonomy candidates or similar-document context.",
            document.pk,
        )
        return empty_taxonomy_candidates(), assigned, ""

    return candidates, assigned, "\n\n".join(context_blocks)


def parse_ai_response(raw: dict) -> ClassificationSuggestions:
    """``raw`` is AIClient.run_llm_query()'s return value - already a
    DocumentClassifierSchema.model_dump(), so every key below is always
    present with the right shape; this only exists to give the rest of the
    module a named, typed boundary instead of passing the client's bare dict
    straight through everywhere.
    """

    def _choice(value: dict | None) -> TaxonomyChoiceDict:
        value = value or {}
        return TaxonomyChoiceDict(
            existing_ids=value.get("existing_ids", []),
            new_names=value.get("new_names", []),
        )

    return ClassificationSuggestions(
        title=raw.get("title", ""),
        tags=_choice(raw.get("tags")),
        correspondents=_choice(raw.get("correspondents")),
        document_types=_choice(raw.get("document_types")),
        storage_paths=_choice(raw.get("storage_paths")),
        dates=raw.get("dates", []),
    )


def _restrict_to_shown_candidates(
    suggestions: ClassificationSuggestions,
    candidates: TaxonomyCandidates,
) -> ClassificationSuggestions:
    """Drop any existing_id the model returned that was never actually
    offered as a candidate in the prompt. The response schema permits any
    integer, so a hallucinated id could otherwise silently resolve to a
    real, visible, but completely unrelated object - this keeps
    "reused an existing value" a fact about what the model was actually
    shown, not just about what integer it happened to emit. When no
    candidates were shown in a category at all (or the field was omitted
    from the response), every existing_id in that category is dropped;
    new_names is never touched here.
    """

    def _restrict(choice: TaxonomyChoiceDict, shown: set[int]) -> TaxonomyChoiceDict:
        return TaxonomyChoiceDict(
            existing_ids=[i for i in choice["existing_ids"] if i in shown],
            new_names=choice["new_names"],
        )

    return ClassificationSuggestions(
        title=suggestions["title"],
        tags=_restrict(
            suggestions["tags"],
            {c["id"] for c in candidates["tags"]},
        ),
        correspondents=_restrict(
            suggestions["correspondents"],
            {c["id"] for c in candidates["correspondents"]},
        ),
        document_types=_restrict(
            suggestions["document_types"],
            {c["id"] for c in candidates["document_types"]},
        ),
        storage_paths=_restrict(
            suggestions["storage_paths"],
            {c["id"] for c in candidates["storage_paths"]},
        ),
        dates=suggestions["dates"],
    )


def get_ai_document_classification(
    document: Document,
    user: User | None = None,
    output_language: str | None = None,
) -> ClassificationSuggestions:
    ai_config = AIConfig()

    if ai_config.llm_embedding_backend:
        candidates, assigned, context = get_taxonomy_context(document, user)
        prompt = build_prompt_with_rag(
            document,
            ai_config,
            candidates=candidates,
            assigned=assigned,
            context=context,
        )
    else:
        candidates = empty_taxonomy_candidates()
        prompt = build_prompt_without_rag(
            document,
            ai_config,
            candidates=candidates,
            assigned=get_assigned_metadata(document, user),
        )

    client = AIClient()
    # Hand the pooled DB connection back while the (slow) LLM query runs so it
    # is not pinned for the call's duration; see paperless_ai.db and #12976.
    with db_connection_released():
        result = client.run_llm_query(prompt)
        suggestions = _restrict_to_shown_candidates(
            parse_ai_response(result),
            candidates,
        )
        if output_language:
            localized = client.run_llm_query(
                build_localization_prompt(suggestions, output_language),
            )
            localized_suggestions = parse_ai_response(localized)

            def _localized_choice(field: str) -> TaxonomyChoiceDict:
                # existing_ids always come from the ORIGINAL suggestions -
                # never from localized_suggestions, whatever the model echoed
                # back there. This is the concrete fix for the bug this
                # feature exists to close: localization must never be able to
                # corrupt an exact taxonomy match.
                return TaxonomyChoiceDict(
                    existing_ids=suggestions[field]["existing_ids"],
                    new_names=localized_suggestions[field]["new_names"]
                    or suggestions[field]["new_names"],
                )

            suggestions = ClassificationSuggestions(
                title=localized_suggestions["title"] or suggestions["title"],
                tags=_localized_choice("tags"),
                correspondents=suggestions["correspondents"],  # never localized
                document_types=_localized_choice("document_types"),
                storage_paths=_localized_choice("storage_paths"),
                dates=suggestions["dates"],
            )
    return suggestions