From b26b1cc6bae252fadbd180e547b40bc032fd6596 Mon Sep 17 00:00:00 2001 From: Trenton Holmes <797416+stumpylog@users.noreply.github.com> Date: Thu, 11 Jun 2026 19:53:51 -0700 Subject: [PATCH] Enhancement(beta): wire indexing pipeline to the sqlite-vec store Co-Authored-By: Claude Sonnet 4.6 --- src/paperless_ai/indexing.py | 43 +++++++++++++++------- src/paperless_ai/tests/test_ai_indexing.py | 34 ++++++++--------- 2 files changed, 47 insertions(+), 30 deletions(-) diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index 5153c0baa..b9360774d 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -1,4 +1,5 @@ import logging +import shutil from collections.abc import Iterable from contextlib import contextmanager from datetime import timedelta @@ -21,7 +22,7 @@ from paperless_ai.embedding import get_embedding_model if TYPE_CHECKING: from llama_index.core.schema import BaseNode - from paperless_ai.vector_store import PaperlessLanceVectorStore + from paperless_ai.vector_store import PaperlessSqliteVecVectorStore logger = logging.getLogger("paperless_ai.indexing") @@ -63,16 +64,31 @@ def queue_llm_index_update_if_needed(*, rebuild: bool, reason: str) -> bool: return True -def get_vector_store() -> "PaperlessLanceVectorStore": - from paperless_ai.vector_store import PaperlessLanceVectorStore +def get_vector_store() -> "PaperlessSqliteVecVectorStore": + from paperless_ai.vector_store import PaperlessSqliteVecVectorStore settings.LLM_INDEX_DIR.mkdir(parents=True, exist_ok=True) - return PaperlessLanceVectorStore( + return PaperlessSqliteVecVectorStore( uri=str(settings.LLM_INDEX_DIR), table_name=LLM_INDEX_TABLE, ) +def _cleanup_legacy_lance_index() -> bool: + """Delete a LanceDB index left by a pre-sqlite-vec version, if present. + + Beta transition policy: no cross-store conversion; the caller forces a + full rebuild (re-embed) instead. Returns True when leftovers were found. + """ + legacy_table = settings.LLM_INDEX_DIR / f"{LLM_INDEX_TABLE}.lance" + found = legacy_table.exists() + if found: + shutil.rmtree(legacy_table, ignore_errors=True) + # faiss-era metadata file, removed on the same occasion + (settings.LLM_INDEX_DIR / "meta.json").unlink(missing_ok=True) + return found + + @contextmanager def write_store(embed_model_name: str | None = None): """Acquire the write lock and yield the vector store. @@ -84,11 +100,11 @@ def write_store(embed_model_name: str | None = None): Pass ``embed_model_name`` whenever the operation may create the table so the model name is recorded in the schema metadata for future mismatch checks. """ - from paperless_ai.vector_store import PaperlessLanceVectorStore + from paperless_ai.vector_store import PaperlessSqliteVecVectorStore settings.LLM_INDEX_DIR.mkdir(parents=True, exist_ok=True) with FileLock(settings.LLM_INDEX_LOCK): - yield PaperlessLanceVectorStore( + yield PaperlessSqliteVecVectorStore( uri=str(settings.LLM_INDEX_DIR), table_name=LLM_INDEX_TABLE, embed_model_name=embed_model_name, @@ -224,6 +240,11 @@ def update_llm_index( rebuild=False, ) -> str: """Rebuild or incrementally update the LLM index.""" + if _cleanup_legacy_lance_index(): + logger.warning( + "Found a LanceDB index from a previous version; forcing a full rebuild.", + ) + rebuild = True documents = Document.objects.all() no_documents = not documents.exists() @@ -251,7 +272,6 @@ def update_llm_index( with write_store(embed_model_name=model_name) as store: if rebuild or not store.table_exists(): - (settings.LLM_INDEX_DIR / "meta.json").unlink(missing_ok=True) logger.info("Rebuilding LLM index.") store.drop_table() for document in iter_wrapper(documents): @@ -276,9 +296,7 @@ def update_llm_index( else "No changes detected in LLM index." ) - store.ensure_document_id_scalar_index() - store.maybe_create_ann_index() - store.compact(retention_seconds=60 * 60) # 1 hour: safe for in-flight readers + store.compact() return msg @@ -294,13 +312,12 @@ def llm_index_add_or_update_document(document: Document): with write_store(embed_model_name=get_configured_model_name(config)) as store: store.upsert_document(str(document.id), new_nodes) - store.ensure_document_id_scalar_index() def llm_index_compact() -> None: - """Compact the index immediately, clearing all MVCC version history.""" + """Compact the index immediately, rebuilding the table to reclaim space.""" with write_store() as store: - store.compact(retention_seconds=0) + store.compact(force=True) def llm_index_remove_document(document: Document): diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index 31e1f6bc8..9ae791613 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -38,7 +38,7 @@ def test_build_document_node(real_document: Document) -> None: @pytest.mark.django_db def test_build_document_node_sets_ref_doc_id(real_document: Document) -> None: """Every node produced by build_document_node must carry the paperless document id - as its ref_doc_id so that the LanceDB adapter's delete(str(doc.id)) works correctly.""" + as its ref_doc_id so that the vector store's delete(str(doc.id)) works correctly.""" nodes = indexing.build_document_node(real_document) assert len(nodes) > 0, "Expected at least one node" for node in nodes: @@ -256,7 +256,7 @@ def test_update_llm_index_partial_update( store = indexing.get_vector_store() assert store.table_exists(), ( - "Expected the LanceDB table to exist after incremental update" + "Expected the vector store table to exist after incremental update" ) @@ -271,7 +271,7 @@ def test_add_or_update_document_updates_existing_entry( store = indexing.get_vector_store() assert store.table_exists(), ( - "Expected the LanceDB table to exist after add-or-update" + "Expected the vector store table to exist after add-or-update" ) @@ -461,7 +461,7 @@ def test_query_similar_documents_empty_allow_list_fails_closed( class TestUpdateLlmIndexEmptyDocumentSet: - """update_llm_index must clear the LanceDB table when all documents are deleted. + """update_llm_index must clear the vector store table when all documents are deleted. Without this, the stale vectors are never cleared and subsequent similarity searches return phantom hits for document IDs that no longer exist in the DB. @@ -491,7 +491,7 @@ class TestUpdateLlmIndexEmptyDocumentSet: store = indexing.get_vector_store() assert store.table_exists(), ( - "Precondition failed: expected the LanceDB table to exist before deletion" + "Precondition failed: expected the vector store table to exist before deletion" ) # Step 2: delete all documents @@ -505,7 +505,7 @@ class TestUpdateLlmIndexEmptyDocumentSet: # Step 4: the table must be absent (no rows) — phantom vectors gone store2 = indexing.get_vector_store() assert not store2.table_exists(), ( - "Expected the LanceDB table to be absent after rebuilding with no documents" + "Expected the vector store table to be absent after rebuilding with no documents" ) @@ -578,11 +578,11 @@ class TestLlmIndexAddOrUpdateDocumentEmptyContent: @pytest.mark.django_db -def test_llm_index_compact_uses_zero_retention( +def test_llm_index_compact_uses_force( temp_llm_index_dir: Path, mocker: pytest_mock.MockerFixture, ) -> None: - """compact must use retention_seconds=0 to clear all MVCC history immediately.""" + """compact must use force=True to rebuild the table and reclaim space immediately.""" mock_store = mocker.MagicMock() mocker.patch( "paperless_ai.indexing.write_store", @@ -594,7 +594,7 @@ def test_llm_index_compact_uses_zero_retention( indexing.llm_index_compact() - mock_store.compact.assert_called_once_with(retention_seconds=0) + mock_store.compact.assert_called_once_with(force=True) @pytest.mark.django_db @@ -678,16 +678,16 @@ class TestLlmIndexLocking: @pytest.mark.django_db @pytest.mark.django_db -class TestLanceDbIndexing: +class TestVectorStoreIndexing: def test_get_vector_store_roundtrip( self, temp_llm_index_dir: Path, mock_embed_model: FakeEmbedding, ) -> None: - from paperless_ai.vector_store import PaperlessLanceVectorStore + from paperless_ai.vector_store import PaperlessSqliteVecVectorStore store = indexing.get_vector_store() - assert isinstance(store, PaperlessLanceVectorStore) + assert isinstance(store, PaperlessSqliteVecVectorStore) def test_add_then_remove_document( self, @@ -697,11 +697,11 @@ class TestLanceDbIndexing: ) -> None: indexing.llm_index_add_or_update_document(real_document) store = indexing.get_vector_store() - table = store.client.open_table(indexing.LLM_INDEX_TABLE) - assert table.count_rows() >= 1 + assert store.table_exists() + assert store.client.execute("SELECT count(*) FROM documents").fetchone()[0] >= 1 indexing.llm_index_remove_document(real_document) - assert store.client.open_table(indexing.LLM_INDEX_TABLE).count_rows() == 0 + assert store.client.execute("SELECT count(*) FROM documents").fetchone()[0] == 0 def test_update_shrinks_chunks_without_orphans( self, @@ -713,13 +713,13 @@ class TestLanceDbIndexing: real_document.save() indexing.llm_index_add_or_update_document(real_document) store = indexing.get_vector_store() - big = store.client.open_table(indexing.LLM_INDEX_TABLE).count_rows() + big = store.client.execute("SELECT count(*) FROM documents").fetchone()[0] real_document.content = "short" # one chunk real_document.save() indexing.llm_index_add_or_update_document(real_document) - rows = store.client.open_table(indexing.LLM_INDEX_TABLE).count_rows() + rows = store.client.execute("SELECT count(*) FROM documents").fetchone()[0] assert rows < big assert rows >= 1