Perf: prefetch notes and custom fields for LLM index text building (#13350)

build_llm_index_text queried Note and CustomFieldInstance (plus its
field FK) per document, uncovered by the earlier correspondent/type/
storage_path/tags prefetch fix. Add notes and custom_fields__field to
the rebuild and scoped document querysets, and read from doc.notes.all()
instead of a fresh Note.objects.filter() so the prefetch is actually used.

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

src/paperless_ai/embedding.py

@@ -7,7 +7,6 @@ if TYPE_CHECKING:
     from llama_index.core.base.embeddings.base import BaseEmbedding
 
 from documents.models import Document
-from documents.models import Note
 from paperless.config import AIConfig
 from paperless.models import LLMEmbeddingBackend
 from paperless.network import PinnedHostAsyncHTTPTransport
@@ -126,7 +125,7 @@ def build_llm_index_text(doc: Document) -> str:
     # prepend. Notes and Custom Fields stay in the body: Notes can be long free
     # text, Custom Fields are dynamic in count and best kept in the embedding.
     lines = [
-        f"Notes: {','.join([str(c.note) for c in Note.objects.filter(document=doc)])}",
+        f"Notes: {','.join([str(c.note) for c in doc.notes.all()])}",
     ]
 
     for instance in doc.custom_fields.all():

src/paperless_ai/indexing.py

@@ -357,7 +357,7 @@ def update_llm_index(
         "correspondent",
         "document_type",
         "storage_path",
-    ).prefetch_related("tags")
+    ).prefetch_related("tags", "notes", "custom_fields__field")
     no_documents = not documents.exists()
 
     # Fast exit before touching config: nothing to index and no existing index.

src/paperless_ai/tests/test_ai_indexing.py

@@ -4,13 +4,18 @@ from unittest.mock import patch
 
 import pytest
 import pytest_mock
+from django.db import connection
 from django.test import override_settings
+from django.test.utils import CaptureQueriesContext
 from django.utils import timezone
 from llama_index.core.schema import MetadataMode
 
 from documents.models import Correspondent
+from documents.models import CustomField
+from documents.models import CustomFieldInstance
 from documents.models import Document
 from documents.models import DocumentType
+from documents.models import Note
 from documents.models import PaperlessTask
 from documents.signals import document_consumption_finished
 from documents.signals import document_updated
@@ -296,14 +301,44 @@ def test_update_llm_index_partial_update(
         )
         before = store.get_modified_times()
 
-    # A further edit, scoped via document_ids to just doc3 -- doc2 must be
+    # new doc, also touched by the scoped update below
+    doc4 = DocumentFactory.create(title="Test Document 4", added=timezone.now())
+
+    # A further edit, scoped via document_ids to doc3 + doc4 -- doc2 must be
     # left exactly as it was, proving document_ids restricts the scan
     # instead of falling back to the whole library.
     doc3.modified = timezone.now()
     doc3.save()
+    doc4.modified = timezone.now()
+    doc4.save()
 
-    result = indexing.update_llm_index(rebuild=False, document_ids=[doc3.pk])
+    # Give both scoped documents a note and a custom field: build_llm_index_text
+    # reads both per document, so without the notes/custom_fields__field
+    # prefetch on scoped_documents, each additional document adds 3 more
+    # queries (N+1 regression) instead of the query count staying flat.
+    custom_field = CustomField.objects.create(
+        name="Priority",
+        data_type=CustomField.FieldDataType.STRING,
+    )
+    for doc in (doc3, doc4):
+        Note.objects.create(document=doc, note=f"a note on {doc.title}")
+        CustomFieldInstance.objects.create(
+            document=doc,
+            field=custom_field,
+            value_text="high",
+        )
+
+    with CaptureQueriesContext(connection) as ctx:
+        result = indexing.update_llm_index(
+            rebuild=False,
+            document_ids=[doc3.pk, doc4.pk],
+        )
     assert result == "LLM index updated successfully."
+    # Notes/custom fields are prefetched in one batch query each (plus one
+    # more for custom_fields__field), not re-queried per document -- an N+1
+    # regression here would scale with document count instead of staying flat
+    # (7 with the prefetch vs. 10 without it, for these 2 documents).
+    assert len(ctx.captured_queries) <= 8
 
     with indexing.get_vector_store() as store:
         after = store.get_modified_times()

src/paperless_ai/tests/test_embedding.py

@@ -54,6 +54,7 @@ def mock_document():
     cf2.field.name = "Field2"
     cf2.value = "Value2"
     doc.custom_fields.all = MagicMock(return_value=[cf1, cf2])
+    doc.notes.all = MagicMock(return_value=[])
 
     return doc
 
@@ -219,28 +220,26 @@ def test_get_configured_model_name_explicit_overrides_default(mock_ai_config):
 
 
 def test_build_llm_index_text(mock_document):
-    with patch("documents.models.Note.objects.filter") as mock_notes_filter:
-        mock_notes_filter.return_value = [
-            MagicMock(note="Note1"),
-            MagicMock(note="Note2"),
-        ]
+    mock_document.notes.all = MagicMock(
+        return_value=[MagicMock(note="Note1"), MagicMock(note="Note2")],
+    )
 
-        result = build_llm_index_text(mock_document)
+    result = build_llm_index_text(mock_document)
 
-        # Structured fields live in node.metadata for LLM context -- not body text
-        assert "Title: Test Title" not in result
-        assert "Created: 2023-01-01" not in result
-        assert "Tags: Tag1, Tag2" not in result
-        assert "Document Type: Invoice" not in result
-        assert "Correspondent: Test Correspondent" not in result
-        assert "Filename:" not in result
-        assert "Storage Path:" not in result
-        assert "Archive Serial Number:" not in result
+    # Structured fields live in node.metadata for LLM context -- not body text
+    assert "Title: Test Title" not in result
+    assert "Created: 2023-01-01" not in result
+    assert "Tags: Tag1, Tag2" not in result
+    assert "Document Type: Invoice" not in result
+    assert "Correspondent: Test Correspondent" not in result
+    assert "Filename:" not in result
+    assert "Storage Path:" not in result
+    assert "Archive Serial Number:" not in result
 
-        # Fields without a metadata equivalent stay in body text
-        assert "Notes: Note1,Note2" in result
-        assert "Content:\n\nThis is the document content." in result
-        assert "Custom Field - Field1: Value1\nCustom Field - Field2: Value2" in result
+    # Fields without a metadata equivalent stay in body text
+    assert "Notes: Note1,Note2" in result
+    assert "Content:\n\nThis is the document content." in result
+    assert "Custom Field - Field1: Value1\nCustom Field - Field2: Value2" in result
 
 
 def test_build_llm_index_text_normalizes_ocr_punctuation_runs(mock_document):
@@ -250,8 +249,7 @@ def test_build_llm_index_text_normalizes_ocr_punctuation_runs(mock_document):
         "Keep short punctuation like INV-100 and ellipses..."
     )
 
-    with patch("documents.models.Note.objects.filter", return_value=[]):
-        result = build_llm_index_text(mock_document)
+    result = build_llm_index_text(mock_document)
 
     assert "Introduction 7" in result
     assert "Hardware Limitation 9" in result