From 0e98a7f1ce102fb7c5f36ef55c134b15b6af4870 Mon Sep 17 00:00:00 2001 From: Trenton H <797416+stumpylog@users.noreply.github.com> Date: Mon, 27 Jul 2026 10:23:43 -0700 Subject: [PATCH] Performance: Scope llm index updates to actually modified documents (#13322) * Perf: scope bulk_update_documents' LLM index refresh to the edited document ids instead of the whole library * Perf: select_related/prefetch_related for the scoped LLM index batch * Perf: select_related/prefetch_related for the full LLM index rebuild path * Fix: address Copilot review findings on LLM index scoping --- src/documents/tasks.py | 2 ++ src/documents/tests/test_tasks.py | 8 ++++++-- src/paperless_ai/indexing.py | 23 ++++++++++++++++++--- src/paperless_ai/tests/test_ai_indexing.py | 24 ++++++++++++++++++++++ 4 files changed, 52 insertions(+), 5 deletions(-) diff --git a/src/documents/tasks.py b/src/documents/tasks.py index 7961af381..a5da8f1cf 100644 --- a/src/documents/tasks.py +++ b/src/documents/tasks.py @@ -311,6 +311,7 @@ def sanity_check(*, raise_on_error: bool = True) -> str: def bulk_update_documents(document_ids) -> None: from documents.search import get_backend + document_ids = list(document_ids) documents = Document.objects.filter(id__in=document_ids) for doc in documents: @@ -331,6 +332,7 @@ def bulk_update_documents(document_ids) -> None: if ai_config.llm_index_enabled: update_llm_index( rebuild=False, + document_ids=document_ids, ) diff --git a/src/documents/tests/test_tasks.py b/src/documents/tests/test_tasks.py index ca57219ad..0e3d3f2a5 100644 --- a/src/documents/tests/test_tasks.py +++ b/src/documents/tests/test_tasks.py @@ -401,6 +401,10 @@ class TestAIIndex(DirectoriesMixin, TestCase): "documents.tasks.update_llm_index", ) as update_llm_index, ): - tasks.bulk_update_documents([doc.pk for doc in docs]) + doc_ids = [doc.pk for doc in docs] + tasks.bulk_update_documents(doc_ids) self.assertEqual(update_document_in_llm_index.apply_async.call_count, 0) - update_llm_index.assert_called_once() + update_llm_index.assert_called_once_with( + rebuild=False, + document_ids=doc_ids, + ) diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index 16aeda776..b413e67c6 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -328,8 +328,16 @@ def update_llm_index( *, iter_wrapper: IterWrapper[Document] = identity, rebuild=False, + document_ids: Iterable[int] | None = None, ) -> str: - """Rebuild or incrementally update the LLM index.""" + """Rebuild or incrementally update the LLM index. + + ``document_ids``, when given, scopes an incremental update to just those + documents instead of scanning the whole library -- callers that already + know which documents changed (e.g. a bulk edit) should pass this to avoid + an O(library size) scan per call. Ignored whenever a rebuild actually + happens, since a rebuild always covers the whole library regardless. + """ with write_store() as store: try: with _exclude_readers(): @@ -345,7 +353,11 @@ def update_llm_index( "LLM index migration requires re-embedding; forcing rebuild.", ) rebuild = True - documents = Document.objects.all() + documents = Document.objects.select_related( + "correspondent", + "document_type", + "storage_path", + ).prefetch_related("tags") no_documents = not documents.exists() # Fast exit before touching config: nothing to index and no existing index. @@ -379,9 +391,14 @@ def update_llm_index( store.add(nodes) msg = "LLM index rebuilt successfully." else: + scoped_documents = ( + documents.filter(id__in=document_ids) + if document_ids is not None + else documents + ) existing = store.get_modified_times() changed = 0 - for document in iter_wrapper(documents): + for document in iter_wrapper(scoped_documents): doc_id = str(document.id) if existing.get(doc_id) == document.modified.isoformat(): continue diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index e69834810..81b4c889c 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -197,6 +197,8 @@ def test_update_llm_index( mock_queryset = MagicMock() mock_queryset.exists.return_value = True mock_queryset.__iter__.return_value = iter([real_document]) + mock_queryset.select_related.return_value = mock_queryset + mock_queryset.prefetch_related.return_value = mock_queryset mock_all.return_value = mock_queryset build_document_node.return_value = [] indexing.update_llm_index(rebuild=True) @@ -216,6 +218,8 @@ def test_update_llm_index_rebuilds_on_model_name_change( mock_queryset = MagicMock() mock_queryset.exists.return_value = True mock_queryset.__iter__.return_value = iter([real_document]) + mock_queryset.select_related.return_value = mock_queryset + mock_queryset.prefetch_related.return_value = mock_queryset mock_all.return_value = mock_queryset with patch( "paperless_ai.indexing.get_configured_model_name", @@ -228,6 +232,8 @@ def test_update_llm_index_rebuilds_on_model_name_change( mock_queryset = MagicMock() mock_queryset.exists.return_value = True mock_queryset.__iter__.return_value = iter([real_document]) + mock_queryset.select_related.return_value = mock_queryset + mock_queryset.prefetch_related.return_value = mock_queryset mock_all.return_value = mock_queryset with patch( "paperless_ai.indexing.get_configured_model_name", @@ -258,6 +264,8 @@ def test_update_llm_index_partial_update( mock_queryset = MagicMock() mock_queryset.exists.return_value = True mock_queryset.__iter__.return_value = iter([real_document, doc2]) + mock_queryset.select_related.return_value = mock_queryset + mock_queryset.prefetch_related.return_value = mock_queryset mock_all.return_value = mock_queryset indexing.update_llm_index(rebuild=True) @@ -286,6 +294,22 @@ def test_update_llm_index_partial_update( assert store.table_exists(), ( "Expected the vector store table to exist after incremental update" ) + before = store.get_modified_times() + + # A further edit, scoped via document_ids to just doc3 -- 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() + + result = indexing.update_llm_index(rebuild=False, document_ids=[doc3.pk]) + assert result == "LLM index updated successfully." + + with indexing.get_vector_store() as store: + after = store.get_modified_times() + + assert after[str(doc3.pk)] == doc3.modified.isoformat() + assert after[str(doc2.pk)] == before[str(doc2.pk)] @pytest.mark.django_db