From 686ff7aa682a3ec06b06c94dafd99ae67c3d3106 Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:16:08 -0700 Subject: [PATCH] Perf: dedupe table_exists() lookups, atomic insert counter, fewer connections in update_llm_index() --- src/paperless_ai/indexing.py | 9 +++-- src/paperless_ai/tables.py | 17 +++++++-- src/paperless_ai/tests/test_ai_indexing.py | 39 +++++++++++++++++++++ src/paperless_ai/tests/test_tables.py | 34 ++++++++++++++++++ src/paperless_ai/tests/test_vector_store.py | 25 +++++++++++++ src/paperless_ai/vector_store.py | 20 +++++++---- 6 files changed, 131 insertions(+), 13 deletions(-) diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index 32c5eb2ae..9678f8fc0 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -393,12 +393,11 @@ def update_llm_index( config = AIConfig() model_name = get_configured_model_name(config) - if not rebuild and llm_index_exists(): + if not rebuild: with read_store() as store: - config_mismatch = store.config_mismatch(model_name) - if config_mismatch: - logger.warning("Embedding model changed; forcing LLM index rebuild.") - rebuild = True + if store.table_exists() and store.config_mismatch(model_name): + logger.warning("Embedding model changed; forcing LLM index rebuild.") + rebuild = True if no_documents: logger.warning("No documents found to index.") diff --git a/src/paperless_ai/tables.py b/src/paperless_ai/tables.py index bc8e962a0..850c0f862 100644 --- a/src/paperless_ai/tables.py +++ b/src/paperless_ai/tables.py @@ -216,8 +216,21 @@ class IndexMetaTable: @staticmethod def increment_total_inserts(conn: sqlite3.Connection, count: int) -> None: - current = IndexMetaTable.get_total_inserts(conn) - IndexMetaTable._set(conn, "total_inserts", str(current + count)) + """Atomically add ``count`` to the stored counter in one statement + (INSERT .. ON CONFLICT DO UPDATE with arithmetic), instead of a + separate read-then-write -- called once per add()/upsert_document(), + so halving the statement count here is a real, if small, per-call + saving. index_meta.value has TEXT affinity, so the incremented + result is stored as its text representation -- get_total_inserts() + already expects that (int(value)), so this is not a behavior + change, only fewer statements. + """ + conn.execute( + "INSERT INTO index_meta (key, value) VALUES ('total_inserts', ?) " + "ON CONFLICT(key) DO UPDATE SET value = " + "CAST(index_meta.value AS INTEGER) + CAST(excluded.value AS INTEGER)", + (str(count),), + ) @staticmethod def reset_total_inserts(conn: sqlite3.Connection, count: int) -> None: diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index 45c7bc3b6..25d0a7204 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -253,6 +253,45 @@ def test_update_llm_index_rebuilds_on_model_name_change( assert store.stored_model_name() == "model-b" +@pytest.mark.django_db +def test_update_llm_index_merges_exists_and_config_mismatch_reads( + temp_llm_index_dir: Path, + real_document: Document, + mock_embed_model: FakeEmbedding, +) -> None: + # Build an initial index so the second call's table_exists()/ + # config_mismatch() checks have something real to check against. + with patch("documents.models.Document.objects.all") as mock_all: + 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 + indexing.update_llm_index(rebuild=True) + + with patch("documents.models.Document.objects.all") as mock_all: + 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.read_store", + wraps=indexing.read_store, + ) as read_store_spy: + indexing.update_llm_index(rebuild=False) + + # Documents exist, so the fast-exit check's `no_documents and ...` + # short-circuits before ever calling llm_index_exists() -- the only + # read_store() call left in this path is the merged table_exists()/ + # config_mismatch() check. Before this task's fix, that merged check + # was two separate read_store() calls (one inside llm_index_exists(), + # one for config_mismatch() right after) -- so this asserts 1, not 2. + assert read_store_spy.call_count == 1 + + @pytest.mark.django_db def test_update_llm_index_partial_update( temp_llm_index_dir: Path, diff --git a/src/paperless_ai/tests/test_tables.py b/src/paperless_ai/tests/test_tables.py index 51cf79b93..1feda711c 100644 --- a/src/paperless_ai/tests/test_tables.py +++ b/src/paperless_ai/tests/test_tables.py @@ -2,6 +2,7 @@ import sqlite3 from collections.abc import Generator import pytest +from pytest_mock import MockerFixture from paperless_ai.tables import ChunkRow from paperless_ai.tables import DocumentChunksTable @@ -288,6 +289,39 @@ class TestIndexMetaTable: IndexMetaTable.increment_total_inserts(conn, 3) assert IndexMetaTable.get_total_inserts(conn) == 8 + def test_increment_total_inserts_is_a_single_statement( + self, + mocker: MockerFixture, + ) -> None: + """ + GIVEN: + - An empty index_meta table + WHEN: + - increment_total_inserts() is called + THEN: + - Exactly one conn.execute() call is made (a single INSERT ... + ON CONFLICT DO UPDATE, not a separate read then write) + """ + + # sqlite3.Connection is an immutable C extension type with no + # instance __dict__, so mocker.spy(conn, "execute") can't shadow + # "execute" on a plain connection ("attribute 'execute' is + # read-only"). A trivial Python subclass gets a normal instance + # __dict__, making the instance spyable while still being a real, + # usable sqlite3.Connection. + class _SpyableConnection(sqlite3.Connection): + pass + + conn = sqlite3.connect(":memory:", factory=_SpyableConnection) + try: + conn.row_factory = sqlite3.Row + IndexMetaTable.create(conn) + execute_spy = mocker.spy(conn, "execute") + IndexMetaTable.increment_total_inserts(conn, 5) + assert execute_spy.call_count == 1 + finally: + conn.close() + def test_reset_total_inserts_sets_absolute_value( self, conn: sqlite3.Connection, diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index 8a08c1ca4..a7defa58d 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -234,6 +234,31 @@ class TestCrud: == 0 ) + def test_upsert_document_checks_table_exists_once( + self, + store: PaperlessSqliteVecVectorStore, + mocker: MockerFixture, + ) -> None: + """ + GIVEN: + - An existing store with one document already indexed + WHEN: + - upsert_document() replaces that document's chunks + THEN: + - table_exists() is queried at most once per call, not twice + (previously: once via _ensure_table(), once via the separate + `if self.table_exists():` delete-chunks guard) + """ + store.add([make_node("a1", 1)]) + # store is a pydantic model, whose __setattr__/__delattr__ reject + # arbitrary instance attributes ("object has no attribute + # 'table_exists'"), so mocker.spy(store, "table_exists") can't + # shadow the method on the instance. Spying on the class works + # (bound method lookup on the instance still resolves through it). + exists_spy = mocker.spy(PaperlessSqliteVecVectorStore, "table_exists") + store.upsert_document(1, [make_node("a2", 1)]) + assert exists_spy.call_count == 1 + class TestBuildWhere: def test_ne_filter_translates_to_not_equal_clause(self) -> None: diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index 79ce4c320..8b2e2881a 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -306,8 +306,8 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): if self._embed_model_name: IndexMetaTable.set_embed_model(self._conn, self._embed_model_name) - def _ensure_table(self, dim: int) -> None: - if not self.table_exists(): + def _ensure_table(self, dim: int, *, table_exists: bool) -> None: + if not table_exists: self._create_table(dim) def _row(self, node: BaseNode) -> _Row: @@ -381,7 +381,10 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): return [] rows = [self._row(node) for node in nodes] with self._transaction(): - self._ensure_table(len(nodes[0].get_embedding())) + self._ensure_table( + len(nodes[0].get_embedding()), + table_exists=self.table_exists(), + ) self._conn.executemany(_INSERT, _vec0_params(rows)) self._index_chunks(rows) self._increment_total_inserts(len(rows)) @@ -402,9 +405,14 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): doc_id = int(document_id) rows = [self._row(node) for node in nodes] with self._transaction(): - if nodes: - self._ensure_table(len(nodes[0].get_embedding())) - if self.table_exists(): + table_exists = self.table_exists() + if nodes and not table_exists: + self._ensure_table( + len(nodes[0].get_embedding()), + table_exists=False, + ) + table_exists = True + if table_exists: self._delete_chunks_by_document_id(doc_id) if rows: self._conn.executemany(_INSERT, _vec0_params(rows))