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

@@ -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,
         )
 
 

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,
+            )

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

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