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)