From 200e1031c962bda290928c65812a5875a5a5e2bd Mon Sep 17 00:00:00 2001 From: Trenton Holmes <797416+stumpylog@users.noreply.github.com> Date: Sun, 26 Jul 2026 15:32:02 -0700 Subject: [PATCH] Perf: select_related/prefetch_related for the full LLM index rebuild path --- src/paperless_ai/indexing.py | 7 ++++++- src/paperless_ai/tests/test_ai_indexing.py | 8 ++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index 5eede29dc..a741c8501 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -381,7 +381,12 @@ def update_llm_index( if rebuild or not store.table_exists(): logger.info("Rebuilding LLM index.") store.drop_table() - for document in iter_wrapper(documents): + rebuild_documents = documents.select_related( + "correspondent", + "document_type", + "storage_path", + ).prefetch_related("tags") + for document in iter_wrapper(rebuild_documents): nodes = build_document_node(document, chunk_size=chunk_size) _embed_nodes(nodes, embed_model) store.add(nodes) diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index 1a14aec4e..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)