diff --git a/src/documents/tasks.py b/src/documents/tasks.py index 7961af381..cf86985dc 100644 --- a/src/documents/tasks.py +++ b/src/documents/tasks.py @@ -331,6 +331,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..c8a8a8f65 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(): @@ -379,9 +387,14 @@ def update_llm_index( store.add(nodes) msg = "LLM index rebuilt successfully." else: + scoped_documents = ( + Document.objects.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..1a14aec4e 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -286,6 +286,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