diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index ef7379462..2bebd0962 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -227,6 +227,12 @@ def update_llm_index( rebuild=False, ) -> str: """Rebuild or incrementally update the LLM index.""" + with write_store() as store: + if store.check_and_run_migrations(): + logger.warning( + "LLM index migration requires re-embedding; forcing rebuild.", + ) + rebuild = True documents = Document.objects.all() no_documents = not documents.exists() diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index e6d9b5534..d83d0f643 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -1,9 +1,14 @@ +import sqlite3 from pathlib import Path import pytest from llama_index.core.schema import TextNode from paperless_ai.vector_store import DB_FILENAME +from paperless_ai.vector_store import DEFAULT_TABLE_NAME +from paperless_ai.vector_store import MIGRATIONS +from paperless_ai.vector_store import SCHEMA_VERSION +from paperless_ai.vector_store import Migration from paperless_ai.vector_store import PaperlessSqliteVecVectorStore DIM = 16 @@ -303,3 +308,123 @@ class TestDbFile: assert ( store.client.execute("PRAGMA journal_mode").fetchone()[0].lower() == "wal" ) + + +class TestMigrations: + """Tests for the schema migration machinery.""" + + def _schema_version(self, store: PaperlessSqliteVecVectorStore) -> int | None: + row = store.client.execute( + "SELECT value FROM index_meta WHERE key = 'schema_version'", + ).fetchone() + return int(row[0]) if row else None + + def test_new_table_records_schema_version(self, store) -> None: + store.add([make_node("a1", "1")]) + assert self._schema_version(store) == SCHEMA_VERSION + + def test_check_migrations_no_table_returns_false(self, store) -> None: + assert store.check_and_run_migrations() is False + + def test_check_migrations_current_version_returns_false(self, store) -> None: + store.add([make_node("a1", "1")]) + assert store.check_and_run_migrations() is False + + def test_reembed_migration_returns_true(self, store, tmp_path: Path) -> None: + store.add([make_node("a1", "1")]) + migration = Migration( + from_version=1, + to_version=2, + kind="re-embed", + description="test re-embed", + ) + MIGRATIONS.append(migration) + try: + from paperless_ai import vector_store as vs_mod + + original = vs_mod.SCHEMA_VERSION + vs_mod.SCHEMA_VERSION = 2 + result = store.check_and_run_migrations() + finally: + MIGRATIONS.remove(migration) + vs_mod.SCHEMA_VERSION = original + assert result is True + + def test_structural_migration_copies_rows_and_updates_version( + self, + store, + tmp_path: Path, + ) -> None: + store.add([make_node("a1", "1"), make_node("b1", "2")]) + + def apply( + src: sqlite3.Connection, + dst: sqlite3.Connection, + dim: int, + ) -> None: + dst.execute( # nosemgrep + f"CREATE VIRTUAL TABLE {DEFAULT_TABLE_NAME} USING vec0(" + "id TEXT PRIMARY KEY, document_id TEXT, modified TEXT," + f" +node_content TEXT, embedding float[{dim}] distance_metric=cosine" + ")", + ) + dst.execute( + "INSERT INTO index_meta (key, value) VALUES ('dim', ?) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value", + (str(dim),), + ) + rows = src.execute( + "SELECT id, document_id, modified, node_content, embedding " + f"FROM {DEFAULT_TABLE_NAME}", + ).fetchall() + dst.execute("BEGIN IMMEDIATE") + dst.executemany( + f"INSERT INTO {DEFAULT_TABLE_NAME} " + "(id, document_id, modified, node_content, embedding) " + "VALUES (?, ?, ?, ?, ?)", + [ + ( + r["id"], + r["document_id"], + r["modified"], + r["node_content"], + bytes(r["embedding"]), + ) + for r in rows + ], + ) + dst.execute( + "INSERT INTO index_meta (key, value) VALUES ('total_inserts', ?) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value", + (str(len(rows)),), + ) + dst.execute("COMMIT") + + migration = Migration( + from_version=1, + to_version=2, + kind="structural", + description="test structural", + apply=apply, + ) + MIGRATIONS.append(migration) + try: + from paperless_ai import vector_store as vs_mod + + original = vs_mod.SCHEMA_VERSION + vs_mod.SCHEMA_VERSION = 2 + result = store.check_and_run_migrations() + finally: + MIGRATIONS.remove(migration) + vs_mod.SCHEMA_VERSION = original + + assert result is False + assert self._schema_version(store) == 2 + ids = {n.node_id for n in store.get_nodes()} + assert ids == {"a1", "b1"} + + def test_compact_preserves_schema_version(self, store) -> None: + store.add([make_node("a1", "1")]) + assert self._schema_version(store) == SCHEMA_VERSION + store.compact(force=True) + assert self._schema_version(store) == SCHEMA_VERSION diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index 9d15c6c7e..84f37503e 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -2,11 +2,15 @@ import json import logging import sqlite3 import struct +from collections.abc import Callable from collections.abc import Iterator from collections.abc import Sequence from contextlib import contextmanager +from dataclasses import dataclass +from dataclasses import field from pathlib import Path from typing import Any +from typing import Literal import sqlite_vec from llama_index.core.bridge.pydantic import PrivateAttr @@ -26,6 +30,11 @@ logger = logging.getLogger("paperless_ai.vector_store") DB_FILENAME = "llmindex.db" DEFAULT_TABLE_NAME = "documents" +# Current schema version. Written to index_meta at table creation and bumped +# whenever a Migration is added to MIGRATIONS. check_and_run_migrations() uses +# this to decide which migrations to run on an existing store. +SCHEMA_VERSION = 1 + # compact(): rebuild when the cumulative rowid count exceeds this multiple of # the live row count. DELETEs on vec0 tables never reclaim space (upstream # asg017/sqlite-vec#54), so per-document re-index churn grows the file until @@ -38,6 +47,38 @@ COMPACT_BLOAT_RATIO = 2.0 _FILTER_COLUMNS = frozenset({"document_id", "modified"}) +@dataclass +class Migration: + """A schema migration for the sqlite-vec vector store. + + kind="structural": rows are copied into a new-schema file with no + re-embedding needed. Supply ``apply(src_conn, dst_conn, dim)`` which + must create the vec0 table in ``dst_conn``, copy all rows from + ``src_conn``, and write ``dim`` / ``embed_model`` / ``total_inserts`` to + ``dst_conn``'s ``index_meta``. ``schema_version`` is written by the + migration runner after ``apply`` returns. + + kind="re-embed": the new schema requires fresh embeddings. + ``check_and_run_migrations()`` returns True when it encounters one of + these so the caller can force a full rebuild (which recreates the table + at the current SCHEMA_VERSION). + """ + + from_version: int + to_version: int + kind: Literal["structural", "re-embed"] + description: str + apply: Callable[[sqlite3.Connection, sqlite3.Connection, int], None] | None = field( + default=None, + repr=False, + ) + + +# Registry of all schema migrations in order. Empty at v1 -- this is the +# baseline. Add entries here (and bump SCHEMA_VERSION) when the schema changes. +MIGRATIONS: list[Migration] = [] + + def _pack(embedding: Sequence[float]) -> bytes: return struct.pack(f"{len(embedding)}f", *embedding) @@ -217,6 +258,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): )""", ) self._meta_set("dim", str(dim)) + self._meta_set("schema_version", str(SCHEMA_VERSION)) if self._embed_model_name: self._meta_set("embed_model", self._embed_model_name) @@ -432,13 +474,14 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): "ON CONFLICT(key) DO UPDATE SET value = excluded.value", ("dim", str(dim)), ) - stored_model = self._meta_get("embed_model") - if stored_model: - new_conn.execute( - "INSERT INTO index_meta (key, value) VALUES (?, ?) " - "ON CONFLICT(key) DO UPDATE SET value = excluded.value", - ("embed_model", stored_model), - ) + for key in ("embed_model", "schema_version"): + value = self._meta_get(key) + if value is not None: + new_conn.execute( + "INSERT INTO index_meta (key, value) VALUES (?, ?) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value", + (key, value), + ) rows = self._conn.execute( "SELECT id, document_id, modified, node_content, embedding " "FROM " + DEFAULT_TABLE_NAME, @@ -482,3 +525,78 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): stale.unlink() Path(compact_path).replace(db_path) self._conn = self._open_connection(db_path) + + def check_and_run_migrations(self) -> bool: + """Apply any pending schema migrations to the store. + + Structural migrations copy live rows into a new-schema file with no + re-embedding. Re-embed migrations cannot be applied automatically; + this method returns True when one is encountered so the caller can + force a full rebuild (which recreates the table at SCHEMA_VERSION). + + Must be called under the write FileLock. No-op when the table does + not exist or is already at SCHEMA_VERSION. + """ + if not self.table_exists(): + return False + + raw = self._meta_get("schema_version") + current = int(raw) if raw is not None else SCHEMA_VERSION + if current >= SCHEMA_VERSION: + return False + + pending = sorted( + [m for m in MIGRATIONS if current <= m.from_version < SCHEMA_VERSION], + key=lambda m: m.from_version, + ) + + for migration in pending: + if migration.kind == "re-embed": + logger.warning( + "LLM index schema v%d -> v%d requires re-embedding (%s); " + "forcing full rebuild.", + migration.from_version, + migration.to_version, + migration.description, + ) + return True + logger.info( + "Running structural LLM index migration v%d -> v%d: %s", + migration.from_version, + migration.to_version, + migration.description, + ) + self._run_structural_migration(migration) + current = migration.to_version + + return False + + def _run_structural_migration(self, migration: Migration) -> None: + """Execute a structural migration using the same file-swap as compact().""" + assert migration.apply is not None, "structural migration must have apply()" + dim = self.vector_dim() + if dim is None: # pragma: no cover + raise RuntimeError("Cannot migrate: no stored vector dimension") + db_path = str(Path(self._uri) / DB_FILENAME) + compact_path = db_path + ".compact" + new_conn = self._open_connection(compact_path) + try: + migration.apply(self._conn, new_conn, dim) + new_conn.execute( + "INSERT INTO index_meta (key, value) VALUES (?, ?) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value", + ("schema_version", str(migration.to_version)), + ) + except BaseException: + new_conn.close() + for p in [compact_path, compact_path + "-wal", compact_path + "-shm"]: + Path(p).unlink(missing_ok=True) + raise + new_conn.close() + self._conn.close() + for suffix in ["-wal", "-shm"]: + stale = Path(compact_path + suffix) + if stale.exists(): + stale.unlink() + Path(compact_path).replace(db_path) + self._conn = self._open_connection(db_path)