diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index 5f90fa051..05e4031bd 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -24,6 +24,7 @@ from paperless_ai.embedding import get_configured_model_name from paperless_ai.embedding import get_embedding_model if TYPE_CHECKING: + from django.db.models import QuerySet from llama_index.core.schema import BaseNode from paperless_ai.vector_store import PaperlessSqliteVecVectorStore @@ -252,6 +253,20 @@ def _safe_related_name(document: Document, field: str) -> str | None: return related.name if related else None +def _document_index_queryset() -> "QuerySet[Document]": + """Document queryset with every relation build_document_node() / + build_llm_index_text() touches -- correspondent, document_type, + storage_path, tags, notes, custom_fields__field -- pre-loaded, so + indexing one document costs a fixed handful of queries regardless of + its tag or custom field count, instead of one query per related object. + """ + return Document.objects.select_related( + "correspondent", + "document_type", + "storage_path", + ).prefetch_related("tags", "notes", "custom_fields__field") + + def build_document_node( document: Document, *, @@ -425,11 +440,7 @@ def update_llm_index( "Skipping LLM index update: migration check deferred; " "will retry next run." ) - documents = Document.objects.select_related( - "correspondent", - "document_type", - "storage_path", - ).prefetch_related("tags", "notes", "custom_fields__field") + documents = _document_index_queryset() no_documents = not documents.exists() # Fast exit before touching config: nothing to index and no existing index. @@ -494,6 +505,21 @@ def update_llm_index( def llm_index_add_or_update_document(document: Document): """Add or atomically replace a document's chunks in the index.""" config = AIConfig() + document_id = document.id + # Re-fetch with the same select_related/prefetch_related shape as the + # bulk path (update_llm_index()) uses: the caller's ``document`` instance + # (e.g. straight off a signal) has none of that loaded, and + # build_document_node()/build_llm_index_text() touch + # correspondent/document_type/storage_path/tags/notes/custom_fields__field + # -- without prefetching, that's one query per related object, including + # one per custom field instance. + document = _document_index_queryset().filter(pk=document_id).first() + if document is None: + logger.info( + "Skipping LLM index update for document %s: it no longer exists.", + document_id, + ) + return new_nodes = build_document_node( document, chunk_size=config.llm_embedding_chunk_size, diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index 2da225053..5663c8a53 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -735,8 +735,71 @@ class TestLlmIndexAddOrUpdateDocumentEmptyContent: ) mock_load = mocker.patch("paperless_ai.indexing.load_or_build_index") - doc = MagicMock(spec=Document) - doc.id = 42 + doc = DocumentFactory.create() + # Must not raise + indexing.llm_index_add_or_update_document(doc) + + mock_load.assert_not_called() + + +@pytest.mark.django_db +class TestLlmIndexAddOrUpdateDocumentPrefetch: + """llm_index_add_or_update_document must prefetch the relations + build_document_node()/build_llm_index_text() touch, not re-query per + tag/note/custom field. + """ + + def test_query_count_does_not_scale_with_custom_field_count( + self, + temp_llm_index_dir: Path, + mock_embed_model: FakeEmbedding, + ) -> None: + """ + GIVEN a document with several custom fields, a note, and a tag + WHEN it is incrementally indexed via llm_index_add_or_update_document + THEN the query count stays flat instead of growing with the number + of custom fields -- a regression here would add one query per + custom field instance (instance.field.name unprefetched), see + build_llm_index_text(). + """ + doc = DocumentFactory.create() + Note.objects.create(document=doc, note="a note") + for i in range(5): + field = CustomField.objects.create( + name=f"Field {i}", + data_type=CustomField.FieldDataType.STRING, + ) + CustomFieldInstance.objects.create( + document=doc, + field=field, + value_text="value", + ) + + with CaptureQueriesContext(connection) as ctx: + indexing.llm_index_add_or_update_document(doc) + + # Flat regardless of custom field count -- an unprefetched + # custom_fields__field would add one query per instance (5 here) on + # top of this budget. + assert len(ctx.captured_queries) <= 10 + + def test_skips_write_when_document_no_longer_exists( + self, + temp_llm_index_dir: Path, + mock_embed_model: FakeEmbedding, + mocker: pytest_mock.MockerFixture, + ) -> None: + """ + GIVEN a document that has been deleted since the caller looked it up + (e.g. a race between a signal firing and its async task running) + WHEN llm_index_add_or_update_document is called with that stale instance + THEN it skips the write instead of raising Document.DoesNotExist + """ + doc = DocumentFactory.create() + doc_id = doc.pk + Document.objects.filter(pk=doc_id).delete() + mock_load = mocker.patch("paperless_ai.indexing.load_or_build_index") + # Must not raise indexing.llm_index_add_or_update_document(doc) @@ -793,8 +856,7 @@ class TestLlmIndexLocking: return_value=[mock_node], ) - doc = MagicMock(spec=Document) - doc.id = 1 + doc = DocumentFactory.create() indexing.llm_index_add_or_update_document(doc) mock_store.upsert_document.assert_called_once() @@ -825,8 +887,7 @@ class TestLlmIndexLocking: return_value=[mock_node], ) - doc = MagicMock(spec=Document) - doc.id = 1 + doc = DocumentFactory.create() indexing.llm_index_add_or_update_document(doc) mock_store.upsert_document.assert_not_called() @@ -863,8 +924,7 @@ class TestLlmIndexLocking: return_value=[mock_node], ) - doc = MagicMock(spec=Document) - doc.id = 1 + doc = DocumentFactory.create() indexing.llm_index_add_or_update_document(doc) mock_store.upsert_document.assert_not_called()