From 98b66fdf24a0cd7ae767e0e91094024b582319cf Mon Sep 17 00:00:00 2001 From: Trenton H <797416+stumpylog@users.noreply.github.com> Date: Mon, 27 Jul 2026 14:50:37 -0700 Subject: [PATCH] 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 --- src/paperless_ai/embedding.py | 3 +- src/paperless_ai/indexing.py | 2 +- src/paperless_ai/tests/test_ai_indexing.py | 39 +++++++++++++++++++-- src/paperless_ai/tests/test_embedding.py | 40 ++++++++++------------ 4 files changed, 58 insertions(+), 26 deletions(-) diff --git a/src/paperless_ai/embedding.py b/src/paperless_ai/embedding.py index 4044d7f08..621b4e979 100644 --- a/src/paperless_ai/embedding.py +++ b/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(): diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index b413e67c6..93b850cbf 100644 --- a/src/paperless_ai/indexing.py +++ b/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. diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index 81b4c889c..667b43d06 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/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() diff --git a/src/paperless_ai/tests/test_embedding.py b/src/paperless_ai/tests/test_embedding.py index fa84d1acc..c5066c8d2 100644 --- a/src/paperless_ai/tests/test_embedding.py +++ b/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