mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-11-05 22:39:23 +00:00
Enhancement(beta): replace LanceDB vector store with sqlite-vec (#12990)
* Chore(beta): add sqlite-vec 0.1.9 dependency Pinned exactly: the 0.1.9 wheels carry no baked SIMD flags (safe on pre-AVX2 CPUs, the point of this migration); the 0.1.10 alphas bake -mavx and would reintroduce the #12970 crash class. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Test(beta): port vector store tests to sqlite-vec backend Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Enhancement(beta): switch AI vector store from LanceDB to sqlite-vec Fixes the non-AVX2 SIGILL class (#12970) at the root: lancedb is no longer imported. sqlite-vec 0.1.9 wheels carry no baked SIMD, vec0 metadata columns give parameterized EQ/IN filtering, WAL preserves the lock-free-reader model, and compact() rebuilds the table because vec0 DELETEs never reclaim space. Implementation notes vs. the Task 3A draft: - compact() uses a file-swap approach (new db file + Path.replace) rather than ALTER TABLE RENAME, which does not cascade to shadow tables in sqlite-vec 0.1.9 (upstream limitation). - Bloat is tracked via a cumulative total_inserts counter in index_meta because the _rowids shadow table does not accumulate deleted rows in 0.1.9 (contrary to the design doc assumption from #54). - None distances from the zero-vector cosine edge case are mapped to similarity 0.0 rather than raising TypeError. - Test suite updated accordingly: _bloat_ratio reads index_meta instead of _rowids; seed collision in force-compact test fixed (seed=100.0). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Enhancement(beta): wire indexing pipeline to the sqlite-vec store Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Enhancement(beta): move filename/storage path/ASN to node metadata Same treatment as title/tags/correspondent in #12944: excluded from the embedded text, visible to the LLM via metadata prepend. Changes embedded text for every document, so it ships inside the sqlite-vec transition, whose forced rebuild re-embeds everything anyway. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Test(beta): cover legacy LanceDB index cleanup and forced rebuild Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Chore(beta): drop lancedb dependency Fixes #12970: the package whose wheels SIGILL on non-AVX2 CPUs is no longer installed at all. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Chore(beta): partial pyrefly cleanup on sqlite-vec vector store - Add MetadataFilter import and isinstance guard in _build_where() - Add query_embedding None guard in query() - Fix dict.get() type-checker ambiguity in get_configured_model_name() Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Chore(beta): drop automatic LanceDB index cleanup on startup Leave legacy Lance directory removal to the user rather than deleting it automatically on first run. Beta policy: user is expected to do a clean re-embed anyway; no need for the system to silently delete their data. Remove _cleanup_legacy_lance_index(), the forced-rebuild path that called it, and the associated tests. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Chore(beta): ruff format pass on sqlite-vec AI files Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Removes the benchmarking file * Try to resolve or silence some semgrep. But we're using SQL here, not an ORM and we control the inputs, not users * Enhancement(beta): add schema migration machinery to sqlite-vec vector store Adds versioned schema migration support modelled after PR #12968's LanceDB approach, adapted for sqlite-vec's file-swap compaction pattern. - SCHEMA_VERSION = 1 written to index_meta at table creation and preserved through compact() - Migration dataclass with from_version, to_version, kind ("structural" or "re-embed"), description, and an optional apply(src, dst, dim) callable - MIGRATIONS registry (empty at v1 baseline); add entries and bump SCHEMA_VERSION when the schema changes - check_and_run_migrations(): structural migrations run via the same file-swap as compact() (no re-embed); re-embed migrations return True so the caller forces a full rebuild - update_llm_index() calls check_and_run_migrations() under the write lock before any indexing work Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Chore(beta): deduplicate vector store internals via helper methods Extract three helpers to remove copy-paste between compact() and _run_structural_migration(): - _meta_set_on(conn, key, value): static upsert into any connection's index_meta; _meta_set() now delegates to it - _create_vec_table(conn, dim): CREATE VIRTUAL TABLE DDL (carries the nosemgrep annotation) - _swap_in_compact(compact_path, db_path): close/replace/reconnect sequence used by both file-swap callers Also normalises compact() error-path cleanup to unlink(missing_ok=True). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Adds equality test and no covers some defensive error handling stuff * Ensures an embed migration stops the migration chain, just in case * Silence one kind right but not really semgrep * Trims dead assignment * Fix(beta): address Copilot review on sqlite-vec vector store Three findings from the PR review: - compact() failure cleanup now unlinks the temporary .compact-wal and .compact-shm files, matching _run_structural_migration(); previously only the main .compact file was removed. - _build_where() fails closed (1 = 0) when filters are requested but none translate, instead of emitting "()" which is invalid SQL; filters scope document access, so an empty translation must match no rows. - Drop the unused table_name constructor parameter (all SQL hardcodes DEFAULT_TABLE_NAME) and its callers in indexing.py. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Enhancement(beta): guard sqlite-vec compaction swap against concurrent readers The compaction/migration file swap replaces the database via os.replace, but the -wal/-shm files are keyed by path, not inode. A reader holding an open connection across the swap leaves the old WAL aliased onto the new file; a subsequent write then corrupts the database (reproduced via PRAGMA integrity_check). Add a cross-process read/write lock (filelock.ReadWriteLock) over the index: - read_store() holds it shared for the whole connection lifetime (and closes the connection on exit); concurrent readers do not block. - compaction and the migration check run under an exclusive lock that drains readers, and skip with an info log on Timeout (maintenance op, retries next run). - Normal writes are untouched: WAL gives reader/writer concurrency and LLM_INDEX_LOCK still serializes writers, so they never block readers. load_or_build_index() now takes the store from the caller's read_store() so the lock and connection span the whole retrieval; chat holds it across the streamed response. Two new settings: LLM_INDEX_RWLOCK and LLM_INDEX_COMPACTION_LOCK_TIMEOUT. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Ensures the store alays cleans up SQLite connections for any operations, even on errors --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Trenton HandClaude Opus 4.8 authored and 
GitHub committed
1 parent
11fb09e4f4
commit
a020f64d08
14 files changed
+1437
-834
No files matched your search
@@ -98,6 +98,13 @@ MODEL_FILE = get_path_from_env(
|
||||
)
|
||||
LLM_INDEX_DIR = DATA_DIR / "llm_index"
|
||||
LLM_INDEX_LOCK = LLM_INDEX_DIR / "index.lock"
|
||||
# Cross-process read/write lock guarding the LLM index compaction/migration
|
||||
# file swap. Readers hold it shared; the swap takes it exclusively so it never
|
||||
# runs while a reader connection is open. Must be a SQLite (.db) file.
|
||||
LLM_INDEX_RWLOCK = LLM_INDEX_DIR / "llmindex.rwlock.db"
|
||||
# Seconds the compaction swap waits for active readers to drain before skipping
|
||||
# this cycle (it is a maintenance operation; the next run retries).
|
||||
LLM_INDEX_COMPACTION_LOCK_TIMEOUT = 30
|
||||
|
||||
LOGGING_DIR = get_path_from_env("PAPERLESS_LOGGING_DIR", DATA_DIR / "log")
|
||||
|
||||
|
||||
+49
-42
@@ -9,6 +9,7 @@ from paperless_ai.db import db_connection_released
|
||||
from paperless_ai.indexing import _document_id_filters
|
||||
from paperless_ai.indexing import get_rag_prompt_helper
|
||||
from paperless_ai.indexing import load_or_build_index
|
||||
from paperless_ai.indexing import read_store
|
||||
|
||||
logger = logging.getLogger("paperless_ai.chat")
|
||||
|
||||
@@ -97,53 +98,59 @@ def _stream_chat_with_documents(query_str: str, documents: list[Document]):
|
||||
from llama_index.core.retrievers import VectorIndexRetriever
|
||||
|
||||
config = AIConfig()
|
||||
index = load_or_build_index(config)
|
||||
filters = _document_id_filters(str(doc.pk) for doc in documents)
|
||||
|
||||
retriever = VectorIndexRetriever(
|
||||
index=index,
|
||||
similarity_top_k=CHAT_RETRIEVER_TOP_K,
|
||||
filters=filters,
|
||||
)
|
||||
# Hold the shared read lock for the whole operation: the query engine
|
||||
# retrieves from the vector store again during synthesis, so the connection
|
||||
# must stay open (and the swap must not run) until the stream finishes.
|
||||
with read_store() as store:
|
||||
index = load_or_build_index(config, store)
|
||||
retriever = VectorIndexRetriever(
|
||||
index=index,
|
||||
similarity_top_k=CHAT_RETRIEVER_TOP_K,
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
# Slow query-embedding + vector search; no Django ORM access happens during
|
||||
# it, so release the pooled DB connection for its duration. See #12976.
|
||||
with db_connection_released():
|
||||
top_nodes = retriever.retrieve(query_str)
|
||||
if not top_nodes:
|
||||
logger.warning("No nodes found for the given documents.")
|
||||
yield CHAT_NO_CONTENT_MESSAGE
|
||||
return
|
||||
# Slow query-embedding + vector search; no Django ORM access happens
|
||||
# during it, so release the pooled DB connection for its duration. See
|
||||
# #12976.
|
||||
with db_connection_released():
|
||||
top_nodes = retriever.retrieve(query_str)
|
||||
if not top_nodes:
|
||||
logger.warning("No nodes found for the given documents.")
|
||||
yield CHAT_NO_CONTENT_MESSAGE
|
||||
return
|
||||
|
||||
client = AIClient()
|
||||
client = AIClient()
|
||||
|
||||
references = _get_document_references(documents, top_nodes)
|
||||
references = _get_document_references(documents, top_nodes)
|
||||
|
||||
prompt_template = PromptTemplate(template=CHAT_PROMPT_TMPL)
|
||||
response_synthesizer = get_response_synthesizer(
|
||||
llm=client.llm,
|
||||
prompt_helper=get_rag_prompt_helper(
|
||||
chunk_size=config.llm_embedding_chunk_size,
|
||||
context_size=config.llm_context_size,
|
||||
),
|
||||
text_qa_template=prompt_template,
|
||||
streaming=True,
|
||||
)
|
||||
query_engine = RetrieverQueryEngine.from_args(
|
||||
retriever=retriever,
|
||||
llm=client.llm,
|
||||
response_synthesizer=response_synthesizer,
|
||||
streaming=True,
|
||||
)
|
||||
prompt_template = PromptTemplate(template=CHAT_PROMPT_TMPL)
|
||||
response_synthesizer = get_response_synthesizer(
|
||||
llm=client.llm,
|
||||
prompt_helper=get_rag_prompt_helper(
|
||||
chunk_size=config.llm_embedding_chunk_size,
|
||||
context_size=config.llm_context_size,
|
||||
),
|
||||
text_qa_template=prompt_template,
|
||||
streaming=True,
|
||||
)
|
||||
query_engine = RetrieverQueryEngine.from_args(
|
||||
retriever=retriever,
|
||||
llm=client.llm,
|
||||
response_synthesizer=response_synthesizer,
|
||||
streaming=True,
|
||||
)
|
||||
|
||||
logger.debug("Document chat query: %s", query_str)
|
||||
# Release the pooled DB connection for the slow streaming LLM response so it
|
||||
# is not pinned for the whole stream; see paperless_ai.db and #12976.
|
||||
with db_connection_released():
|
||||
response_stream = query_engine.query(query_str)
|
||||
for chunk in response_stream.response_gen:
|
||||
yield chunk
|
||||
sys.stdout.flush()
|
||||
logger.debug("Document chat query: %s", query_str)
|
||||
# Release the pooled DB connection for the slow streaming LLM response
|
||||
# so it is not pinned for the whole stream; see paperless_ai.db and
|
||||
# #12976.
|
||||
with db_connection_released():
|
||||
response_stream = query_engine.query(query_str)
|
||||
for chunk in response_stream.response_gen:
|
||||
yield chunk
|
||||
sys.stdout.flush()
|
||||
|
||||
if references:
|
||||
yield _format_chat_metadata_trailer(references)
|
||||
if references:
|
||||
yield _format_chat_metadata_trailer(references)
|
||||
@@ -99,9 +99,13 @@ _DEFAULT_MODEL_NAMES = {
|
||||
|
||||
def get_configured_model_name(config: AIConfig) -> str:
|
||||
"""Return the canonical name of the currently configured embedding model."""
|
||||
default = _DEFAULT_MODEL_NAMES.get(
|
||||
config.llm_embedding_backend,
|
||||
"sentence-transformers/all-MiniLM-L6-v2",
|
||||
# dict.get(key, default) overload resolution fails for TextChoices keys in some
|
||||
# type checkers; use `or` fallback to avoid the ambiguity.
|
||||
default = (
|
||||
_DEFAULT_MODEL_NAMES.get(
|
||||
config.llm_embedding_backend,
|
||||
)
|
||||
or "sentence-transformers/all-MiniLM-L6-v2"
|
||||
)
|
||||
return config.llm_embedding_model or default
|
||||
|
||||
@@ -112,15 +116,11 @@ def _normalize_llm_index_text(text: str) -> str:
|
||||
|
||||
|
||||
def build_llm_index_text(doc: Document) -> str:
|
||||
# TODO: Filename, Storage Path, and Archive Serial Number are short structured
|
||||
# values that could move to node.metadata (excluded from embeddings, visible to
|
||||
# LLM via metadata prepend) — same pattern as title/tags/correspondent. Notes
|
||||
# and Custom Fields should stay here: Notes can be long free text, Custom Fields
|
||||
# are dynamic in count and best kept in the embedding.
|
||||
# Short structured fields (filename, storage path, ASN, title, tags, ...) live
|
||||
# in node.metadata: excluded from embeddings, shown to the LLM via metadata
|
||||
# prepend. Notes and Custom Fields stay in the body: Notes can be long free
|
||||
# text, Custom Fields are dynamic in count and best kept in the embedding.
|
||||
lines = [
|
||||
f"Filename: {doc.filename}",
|
||||
f"Storage Path: {doc.storage_path.name if doc.storage_path else ''}",
|
||||
f"Archive Serial Number: {doc.archive_serial_number or ''}",
|
||||
f"Notes: {','.join([str(c.note) for c in Note.objects.filter(document=doc)])}",
|
||||
]
|
||||
|
||||
|
||||
+149
-45
@@ -7,6 +7,8 @@ from typing import TYPE_CHECKING
|
||||
from django.conf import settings
|
||||
from django.utils import timezone
|
||||
from filelock import FileLock
|
||||
from filelock import ReadWriteLock
|
||||
from filelock import Timeout
|
||||
|
||||
from documents.models import Document
|
||||
from documents.models import PaperlessTask
|
||||
@@ -21,13 +23,11 @@ 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")
|
||||
|
||||
LLM_INDEX_TABLE = "documents"
|
||||
|
||||
RAG_NUM_OUTPUT = 512
|
||||
RAG_CHUNK_OVERLAP = 200
|
||||
|
||||
@@ -63,36 +63,108 @@ 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,
|
||||
)
|
||||
|
||||
|
||||
# --- LLM index locking ---------------------------------------------------
|
||||
#
|
||||
# Two locks guard the index; they answer different questions and are NOT
|
||||
# interchangeable:
|
||||
#
|
||||
# * settings.LLM_INDEX_LOCK (FileLock, exclusive) -- serializes WRITERS against
|
||||
# each other, so only one rebuild/upsert/delete/compaction runs at a time.
|
||||
# Taken by write_store(). Readers never take it, so it never blocks reads.
|
||||
#
|
||||
# * settings.LLM_INDEX_RWLOCK (ReadWriteLock) -- coordinates readers against the
|
||||
# compaction/migration file swap. read_store() takes it SHARED (readers run
|
||||
# concurrently); _exclude_readers() takes it EXCLUSIVE, only for the swap, so
|
||||
# the database file is never replaced while a reader connection is open (that
|
||||
# would alias the old WAL onto the new file and corrupt it).
|
||||
#
|
||||
# | vs another writer | vs a reader
|
||||
# -----------------+-------------------+----------------------------
|
||||
# normal write | LLM_INDEX_LOCK | nothing (WAL gives MVCC)
|
||||
# compaction/swap | LLM_INDEX_LOCK | LLM_INDEX_RWLOCK (exclusive)
|
||||
# reader | nothing (WAL) | LLM_INDEX_RWLOCK (shared)
|
||||
#
|
||||
# They can't be merged into one ReadWriteLock: a normal write must exclude other
|
||||
# writers WITHOUT blocking readers (WAL already gives reader/writer concurrency),
|
||||
# and ReadWriteLock has no "exclusive vs writers, shared vs readers" mode. Only
|
||||
# the swap needs to exclude readers.
|
||||
def _index_rwlock() -> ReadWriteLock:
|
||||
"""Return a fresh read/write lock instance for the index swap.
|
||||
|
||||
``is_singleton=False`` so reads and the swap always coordinate through
|
||||
SQLite (the actual cross-process case) rather than hitting the in-process
|
||||
reentrant-upgrade guard; callers must ``close()`` it (the context managers
|
||||
below do).
|
||||
"""
|
||||
settings.LLM_INDEX_DIR.mkdir(parents=True, exist_ok=True)
|
||||
return ReadWriteLock(str(settings.LLM_INDEX_RWLOCK), is_singleton=False)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def read_store():
|
||||
"""Acquire the shared read lock and yield the vector store for a read.
|
||||
|
||||
The shared lock is held for the whole lifetime of the connection (and
|
||||
closed on exit) so the compaction/migration swap, which takes the exclusive
|
||||
lock, never runs while this connection is open. Concurrent readers do not
|
||||
block each other; only the swap does.
|
||||
"""
|
||||
lock = _index_rwlock()
|
||||
try:
|
||||
with lock.read_lock(), get_vector_store() as store:
|
||||
yield store
|
||||
finally:
|
||||
lock.close()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _exclude_readers():
|
||||
"""Acquire exclusive index access, blocking until readers have drained.
|
||||
|
||||
The exclusive counterpart to ``read_store()``: a compaction or migration
|
||||
must not run while any reader connection is open. Raises
|
||||
:class:`filelock.Timeout` if active readers do not drain within
|
||||
``LLM_INDEX_COMPACTION_LOCK_TIMEOUT``; callers skip the operation on timeout.
|
||||
"""
|
||||
lock = _index_rwlock()
|
||||
try:
|
||||
with lock.write_lock(timeout=settings.LLM_INDEX_COMPACTION_LOCK_TIMEOUT):
|
||||
yield
|
||||
finally:
|
||||
lock.close()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def write_store(embed_model_name: str | None = None):
|
||||
"""Acquire the write lock and yield the vector store.
|
||||
|
||||
All mutating operations (upsert, delete, rebuild, compact) must go through
|
||||
this context manager to serialise concurrent Celery writers.
|
||||
Read paths use ``get_vector_store()`` directly — no lock needed.
|
||||
Read paths use ``read_store()`` so they hold the shared read lock.
|
||||
|
||||
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(
|
||||
with (
|
||||
FileLock(settings.LLM_INDEX_LOCK),
|
||||
PaperlessSqliteVecVectorStore(
|
||||
uri=str(settings.LLM_INDEX_DIR),
|
||||
table_name=LLM_INDEX_TABLE,
|
||||
embed_model_name=embed_model_name,
|
||||
)
|
||||
) as store,
|
||||
):
|
||||
yield store
|
||||
|
||||
|
||||
def build_document_node(
|
||||
@@ -114,6 +186,9 @@ def build_document_node(
|
||||
"document_type": document.document_type.name
|
||||
if document.document_type
|
||||
else None,
|
||||
"filename": document.filename,
|
||||
"storage_path": document.storage_path.name if document.storage_path else None,
|
||||
"archive_serial_number": document.archive_serial_number,
|
||||
"created": document.created.isoformat() if document.created else None,
|
||||
"added": document.added.isoformat() if document.added else None,
|
||||
"modified": document.modified.isoformat(),
|
||||
@@ -140,23 +215,27 @@ def build_document_node(
|
||||
return parser.get_nodes_from_documents([doc])
|
||||
|
||||
|
||||
def load_or_build_index(config: AIConfig):
|
||||
"""Return a VectorStoreIndex backed by the vector store."""
|
||||
def load_or_build_index(config: AIConfig, store: "PaperlessSqliteVecVectorStore"):
|
||||
"""Return a VectorStoreIndex backed by ``store``.
|
||||
|
||||
``store`` is supplied by the caller's ``read_store()`` context so the shared
|
||||
read lock and the connection stay alive for the whole retrieval.
|
||||
"""
|
||||
import llama_index.core.settings as llama_settings
|
||||
from llama_index.core import VectorStoreIndex
|
||||
|
||||
embed_model = get_embedding_model(config)
|
||||
llama_settings.Settings.embed_model = embed_model
|
||||
vector_store = get_vector_store()
|
||||
return VectorStoreIndex.from_vector_store(
|
||||
vector_store=vector_store,
|
||||
vector_store=store,
|
||||
embed_model=embed_model,
|
||||
)
|
||||
|
||||
|
||||
def llm_index_exists() -> bool:
|
||||
"""True when the index table exists on disk."""
|
||||
return get_vector_store().table_exists()
|
||||
with read_store() as store:
|
||||
return store.table_exists()
|
||||
|
||||
|
||||
def get_rag_chunk_size() -> int:
|
||||
@@ -224,6 +303,21 @@ def update_llm_index(
|
||||
rebuild=False,
|
||||
) -> str:
|
||||
"""Rebuild or incrementally update the LLM index."""
|
||||
with write_store() as store:
|
||||
try:
|
||||
with _exclude_readers():
|
||||
needs_reembed = store.check_and_run_migrations()
|
||||
except Timeout:
|
||||
logger.info(
|
||||
"Skipping LLM index migration check: index readers are active; "
|
||||
"will retry next run.",
|
||||
)
|
||||
needs_reembed = False
|
||||
if needs_reembed:
|
||||
logger.warning(
|
||||
"LLM index migration requires re-embedding; forcing rebuild.",
|
||||
)
|
||||
rebuild = True
|
||||
documents = Document.objects.all()
|
||||
no_documents = not documents.exists()
|
||||
|
||||
@@ -235,13 +329,12 @@ def update_llm_index(
|
||||
config = AIConfig()
|
||||
model_name = get_configured_model_name(config)
|
||||
|
||||
if (
|
||||
not rebuild
|
||||
and llm_index_exists()
|
||||
and get_vector_store().config_mismatch(model_name)
|
||||
):
|
||||
logger.warning("Embedding model changed; forcing LLM index rebuild.")
|
||||
rebuild = True
|
||||
if not rebuild and llm_index_exists():
|
||||
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 no_documents:
|
||||
logger.warning("No documents found to index.")
|
||||
@@ -251,7 +344,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 +368,14 @@ 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
|
||||
try:
|
||||
with _exclude_readers():
|
||||
store.compact()
|
||||
except Timeout:
|
||||
logger.info(
|
||||
"Skipping LLM index compaction: index readers are active; "
|
||||
"will retry next run.",
|
||||
)
|
||||
return msg
|
||||
|
||||
|
||||
@@ -294,13 +391,19 @@ 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)
|
||||
try:
|
||||
with _exclude_readers():
|
||||
store.compact(force=True)
|
||||
except Timeout:
|
||||
logger.info(
|
||||
"Skipping LLM index compaction: index readers are active; "
|
||||
"will retry next run.",
|
||||
)
|
||||
|
||||
|
||||
def llm_index_remove_document(document: Document):
|
||||
@@ -367,30 +470,31 @@ def query_similar_documents(
|
||||
|
||||
from llama_index.core.retrievers import VectorIndexRetriever
|
||||
|
||||
index = load_or_build_index(config)
|
||||
|
||||
filters = (
|
||||
_document_id_filters(allowed_document_ids)
|
||||
if allowed_document_ids is not None
|
||||
else None
|
||||
)
|
||||
|
||||
retriever = VectorIndexRetriever(
|
||||
index=index,
|
||||
similarity_top_k=top_k,
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
query_text = truncate_content(
|
||||
(document.title or "") + "\n" + (document.content or ""),
|
||||
chunk_size=config.llm_embedding_chunk_size,
|
||||
context_size=config.llm_context_size,
|
||||
)
|
||||
# The retrieve() call generates a query embedding (a slow external request)
|
||||
# and searches the vector store; no Django ORM access happens during it, so
|
||||
# release the pooled DB connection for its duration. See #12976.
|
||||
with db_connection_released():
|
||||
results = retriever.retrieve(query_text)
|
||||
# Hold the shared read lock for the whole retrieval so the connection is
|
||||
# never open across a compaction swap. The retrieve() call generates a
|
||||
# query embedding (a slow external request) and searches the vector store;
|
||||
# no Django ORM access happens during it, so release the pooled DB
|
||||
# connection for its duration. See #12976.
|
||||
with read_store() as store:
|
||||
index = load_or_build_index(config, store)
|
||||
retriever = VectorIndexRetriever(
|
||||
index=index,
|
||||
similarity_top_k=top_k,
|
||||
filters=filters,
|
||||
)
|
||||
with db_connection_released():
|
||||
results = retriever.retrieve(query_text)
|
||||
|
||||
retrieved_document_ids: list[int] = []
|
||||
for node in results:
|
||||
|
||||
@@ -10,6 +10,7 @@ from pytest_django.fixtures import SettingsWrapper
|
||||
def temp_llm_index_dir(tmp_path: Path, settings: SettingsWrapper) -> Path:
|
||||
settings.LLM_INDEX_DIR = tmp_path
|
||||
settings.LLM_INDEX_LOCK = tmp_path / "index.lock"
|
||||
settings.LLM_INDEX_RWLOCK = tmp_path / "llmindex.rwlock.db"
|
||||
return tmp_path
|
||||
|
||||
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import patch
|
||||
@@ -7,6 +6,7 @@ import pytest
|
||||
import pytest_mock
|
||||
from django.test import override_settings
|
||||
from django.utils import timezone
|
||||
from llama_index.core.schema import MetadataMode
|
||||
|
||||
from documents.models import Document
|
||||
from documents.models import PaperlessTask
|
||||
@@ -17,6 +17,7 @@ from documents.tests.factories import PaperlessTaskFactory
|
||||
from paperless.models import ApplicationConfiguration
|
||||
from paperless_ai import indexing
|
||||
from paperless_ai.tests.conftest import FakeEmbedding
|
||||
from paperless_ai.vector_store import PaperlessSqliteVecVectorStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -33,12 +34,22 @@ def test_build_document_node(real_document: Document) -> None:
|
||||
nodes = indexing.build_document_node(real_document)
|
||||
assert len(nodes) > 0
|
||||
assert nodes[0].metadata["document_id"] == str(real_document.id)
|
||||
assert nodes[0].metadata["filename"] == real_document.filename
|
||||
assert nodes[0].metadata["storage_path"] == (
|
||||
real_document.storage_path.name if real_document.storage_path else None
|
||||
)
|
||||
assert (
|
||||
nodes[0].metadata["archive_serial_number"]
|
||||
== real_document.archive_serial_number
|
||||
)
|
||||
assert "filename" in nodes[0].excluded_embed_metadata_keys
|
||||
assert "filename" not in nodes[0].excluded_llm_metadata_keys
|
||||
|
||||
|
||||
@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:
|
||||
@@ -58,8 +69,6 @@ def test_build_document_node_excludes_metadata_from_embedding(
|
||||
double the token count and exceed embedding models with small context
|
||||
windows (e.g. nomic-embed-text via Ollama defaults to num_ctx=2048).
|
||||
"""
|
||||
from llama_index.core.schema import MetadataMode
|
||||
|
||||
nodes = indexing.build_document_node(real_document)
|
||||
for node in nodes:
|
||||
embed_text = node.get_content(metadata_mode=MetadataMode.EMBED)
|
||||
@@ -91,8 +100,6 @@ def test_build_document_node_excludes_document_id_from_llm_context(
|
||||
real_document: Document,
|
||||
) -> None:
|
||||
"""document_id is an internal key and must not appear in LLM context text."""
|
||||
from llama_index.core.schema import MetadataMode
|
||||
|
||||
nodes = indexing.build_document_node(real_document)
|
||||
assert len(nodes) > 0
|
||||
for node in nodes:
|
||||
@@ -154,29 +161,6 @@ def test_update_llm_index(
|
||||
build_document_node.assert_called_once_with(real_document, chunk_size=512)
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_update_llm_index_cleans_stale_meta_on_rebuild(
|
||||
temp_llm_index_dir: Path,
|
||||
real_document: Document,
|
||||
mock_embed_model: FakeEmbedding,
|
||||
) -> None:
|
||||
# A meta.json left over from the FAISS era (or written by older code) must be
|
||||
# deleted on rebuild so stale artifacts don't accumulate on disk.
|
||||
stale_meta = temp_llm_index_dir / "meta.json"
|
||||
stale_meta.write_text(json.dumps({"embedding_model": "old", "dim": 1}))
|
||||
|
||||
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_all.return_value = mock_queryset
|
||||
indexing.update_llm_index(rebuild=True)
|
||||
|
||||
assert not stale_meta.exists(), (
|
||||
"update_llm_index(rebuild=True) must remove stale meta.json"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_update_llm_index_rebuilds_on_model_name_change(
|
||||
temp_llm_index_dir: Path,
|
||||
@@ -207,10 +191,10 @@ def test_update_llm_index_rebuilds_on_model_name_change(
|
||||
):
|
||||
indexing.update_llm_index(rebuild=False)
|
||||
|
||||
store = indexing.get_vector_store()
|
||||
# Schema metadata only updates when the table is dropped and recreated, never on
|
||||
# incremental writes -- so "model-b" here proves a full rebuild happened.
|
||||
assert store.stored_model_name() == "model-b"
|
||||
with indexing.get_vector_store() as store:
|
||||
# Schema metadata only updates when the table is dropped and recreated, never
|
||||
# on incremental writes -- so "model-b" here proves a full rebuild happened.
|
||||
assert store.stored_model_name() == "model-b"
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
@@ -254,10 +238,10 @@ def test_update_llm_index_partial_update(
|
||||
|
||||
indexing.update_llm_index(rebuild=False)
|
||||
|
||||
store = indexing.get_vector_store()
|
||||
assert store.table_exists(), (
|
||||
"Expected the LanceDB table to exist after incremental update"
|
||||
)
|
||||
with indexing.get_vector_store() as store:
|
||||
assert store.table_exists(), (
|
||||
"Expected the vector store table to exist after incremental update"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
@@ -269,10 +253,10 @@ def test_add_or_update_document_updates_existing_entry(
|
||||
indexing.update_llm_index(rebuild=True)
|
||||
indexing.llm_index_add_or_update_document(real_document)
|
||||
|
||||
store = indexing.get_vector_store()
|
||||
assert store.table_exists(), (
|
||||
"Expected the LanceDB table to exist after add-or-update"
|
||||
)
|
||||
with indexing.get_vector_store() as store:
|
||||
assert store.table_exists(), (
|
||||
"Expected the vector store table to exist after add-or-update"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
@@ -461,7 +445,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.
|
||||
@@ -489,10 +473,11 @@ class TestUpdateLlmIndexEmptyDocumentSet:
|
||||
)
|
||||
indexing.update_llm_index(rebuild=True)
|
||||
|
||||
store = indexing.get_vector_store()
|
||||
assert store.table_exists(), (
|
||||
"Precondition failed: expected the LanceDB table to exist before deletion"
|
||||
)
|
||||
with indexing.get_vector_store() as store:
|
||||
assert store.table_exists(), (
|
||||
"Precondition failed: expected the vector store table to exist "
|
||||
"before deletion"
|
||||
)
|
||||
|
||||
# Step 2: delete all documents
|
||||
Document.objects.all().delete()
|
||||
@@ -503,10 +488,11 @@ class TestUpdateLlmIndexEmptyDocumentSet:
|
||||
indexing.update_llm_index(rebuild=True)
|
||||
|
||||
# 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"
|
||||
)
|
||||
with indexing.get_vector_store() as store2:
|
||||
assert not store2.table_exists(), (
|
||||
"Expected the vector store table to be absent after rebuilding "
|
||||
"with no documents"
|
||||
)
|
||||
|
||||
|
||||
class TestDocumentUpdatedSignalTriggersLlmReindex:
|
||||
@@ -578,11 +564,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 +580,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 +664,14 @@ 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
|
||||
|
||||
store = indexing.get_vector_store()
|
||||
assert isinstance(store, PaperlessLanceVectorStore)
|
||||
with indexing.get_vector_store() as store:
|
||||
assert isinstance(store, PaperlessSqliteVecVectorStore)
|
||||
|
||||
def test_add_then_remove_document(
|
||||
self,
|
||||
@@ -696,12 +680,13 @@ class TestLanceDbIndexing:
|
||||
real_document: Document,
|
||||
) -> 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
|
||||
with indexing.get_vector_store() as store:
|
||||
assert store.table_exists()
|
||||
count_sql = "SELECT count(*) FROM documents"
|
||||
assert store.client.execute(count_sql).fetchone()[0] >= 1
|
||||
|
||||
indexing.llm_index_remove_document(real_document)
|
||||
assert store.client.open_table(indexing.LLM_INDEX_TABLE).count_rows() == 0
|
||||
indexing.llm_index_remove_document(real_document)
|
||||
assert store.client.execute(count_sql).fetchone()[0] == 0
|
||||
|
||||
def test_update_shrinks_chunks_without_orphans(
|
||||
self,
|
||||
@@ -712,16 +697,17 @@ class TestLanceDbIndexing:
|
||||
real_document.content = "word " * 4000 # many chunks
|
||||
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()
|
||||
count_sql = "SELECT count(*) FROM documents"
|
||||
with indexing.get_vector_store() as store:
|
||||
big = store.client.execute(count_sql).fetchone()[0]
|
||||
|
||||
real_document.content = "short" # one chunk
|
||||
real_document.save()
|
||||
indexing.llm_index_add_or_update_document(real_document)
|
||||
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()
|
||||
assert rows < big
|
||||
assert rows >= 1
|
||||
rows = store.client.execute(count_sql).fetchone()[0]
|
||||
assert rows < big
|
||||
assert rows >= 1
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
|
||||
@@ -3,9 +3,13 @@ from unittest.mock import MagicMock
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from llama_index.core import settings as llama_settings
|
||||
from llama_index.core.embeddings.mock_embed_model import MockEmbedding
|
||||
from llama_index.core.schema import TextNode
|
||||
|
||||
from documents.tests.factories import DocumentFactory
|
||||
from paperless_ai import chat
|
||||
from paperless_ai import indexing
|
||||
from paperless_ai.chat import CHAT_ERROR_MESSAGE
|
||||
from paperless_ai.chat import CHAT_METADATA_DELIMITER
|
||||
from paperless_ai.chat import stream_chat_with_documents
|
||||
@@ -13,9 +17,6 @@ from paperless_ai.chat import stream_chat_with_documents
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_embed_model():
|
||||
from llama_index.core import settings as llama_settings
|
||||
from llama_index.core.embeddings.mock_embed_model import MockEmbedding
|
||||
|
||||
# Use a real BaseEmbedding subclass to satisfy llama-index 0.14 validation
|
||||
llama_settings.Settings.embed_model = MockEmbedding(embed_dim=1536)
|
||||
yield
|
||||
@@ -241,8 +242,6 @@ class TestStreamChatRetrieval:
|
||||
temp_llm_index_dir,
|
||||
mock_embed_model,
|
||||
) -> None:
|
||||
from documents.tests.factories import DocumentFactory
|
||||
|
||||
doc = DocumentFactory.create(content="hello world")
|
||||
# Nothing indexed for this document yet.
|
||||
out = list(chat.stream_chat_with_documents("question?", [doc]))
|
||||
@@ -258,9 +257,6 @@ class TestStreamChatRetrieval:
|
||||
requested documents only — content from other indexed documents must
|
||||
not be surfaced.
|
||||
"""
|
||||
from documents.tests.factories import DocumentFactory
|
||||
from paperless_ai import indexing
|
||||
|
||||
included = DocumentFactory.create(content="included document content")
|
||||
excluded = DocumentFactory.create(content="excluded document content")
|
||||
indexing.llm_index_add_or_update_document(included)
|
||||
|
||||
@@ -224,15 +224,17 @@ def test_build_llm_index_text(mock_document):
|
||||
|
||||
result = build_llm_index_text(mock_document)
|
||||
|
||||
# Structured fields live in node.metadata for LLM context — not body text
|
||||
# Structured fields live in node.metadata for LLM context -- not body text
|
||||
assert "Title: Test Title" not in result
|
||||
assert "Created: 2023-01-01" not in result
|
||||
assert "Tags: Tag1, Tag2" not in result
|
||||
assert "Document Type: Invoice" not in result
|
||||
assert "Correspondent: Test Correspondent" not in result
|
||||
assert "Filename:" not in result
|
||||
assert "Storage Path:" not in result
|
||||
assert "Archive Serial Number:" not in result
|
||||
|
||||
# Fields without a metadata equivalent stay in body text
|
||||
assert "Filename: test_file.pdf" in result
|
||||
assert "Notes: Note1,Note2" in result
|
||||
assert "Content:\n\nThis is the document content." in result
|
||||
assert "Custom Field - Field1: Value1\nCustom Field - Field2: Value2" in result
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
import logging
|
||||
import sqlite3
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from django.conf import settings
|
||||
from filelock import ReadWriteLock
|
||||
from llama_index.core.schema import TextNode
|
||||
from pytest_django.fixtures import SettingsWrapper
|
||||
|
||||
from paperless_ai import indexing
|
||||
from paperless_ai.vector_store import PaperlessSqliteVecVectorStore
|
||||
|
||||
DIM = 8
|
||||
|
||||
|
||||
def _node(node_id: str, document_id: str, *, seed: float = 0.0) -> TextNode:
|
||||
node = TextNode(
|
||||
id_=node_id,
|
||||
text="chunk",
|
||||
metadata={"document_id": document_id, "modified": "2026-06-01T00:00:00"},
|
||||
)
|
||||
node.relationships = {}
|
||||
node.embedding = [seed + i / 100 for i in range(DIM)]
|
||||
return node
|
||||
|
||||
|
||||
def _seed_bloated_index(index_dir: Path) -> None:
|
||||
"""Create an index whose cumulative inserts far exceed live rows."""
|
||||
store = PaperlessSqliteVecVectorStore(uri=str(index_dir))
|
||||
store.add([_node(f"d{j}", str(j), seed=float(j)) for j in range(20)])
|
||||
for cycle in range(6):
|
||||
for j in range(20):
|
||||
store.upsert_document(
|
||||
str(j),
|
||||
[_node(f"d{j}-c{cycle}", str(j), seed=float(j))],
|
||||
)
|
||||
store.client.close()
|
||||
|
||||
|
||||
def _bloat_ratio(index_dir: Path) -> float:
|
||||
store = PaperlessSqliteVecVectorStore(uri=str(index_dir))
|
||||
live = store.client.execute("SELECT count(*) FROM documents").fetchone()[0]
|
||||
row = store.client.execute(
|
||||
"SELECT value FROM index_meta WHERE key = 'total_inserts'",
|
||||
).fetchone()
|
||||
total = int(row["value"]) if row else live
|
||||
store.client.close()
|
||||
return total / max(live, 1)
|
||||
|
||||
|
||||
def _integrity_ok(index_dir: Path) -> bool:
|
||||
store = PaperlessSqliteVecVectorStore(uri=str(index_dir))
|
||||
result = store.client.execute("PRAGMA integrity_check").fetchone()[0]
|
||||
rows = store.client.execute("SELECT count(*) FROM documents").fetchone()[0]
|
||||
store.client.close()
|
||||
return result == "ok" and rows == 20
|
||||
|
||||
|
||||
def _reader_lock() -> ReadWriteLock:
|
||||
# A distinct instance simulates a reader in another process: it coordinates
|
||||
# with the production lock purely through SQLite, never reentrant upgrade.
|
||||
return ReadWriteLock(str(settings.LLM_INDEX_RWLOCK), is_singleton=False)
|
||||
|
||||
|
||||
class TestCompactionLock:
|
||||
def test_compaction_skips_when_a_reader_holds_the_lock(
|
||||
self,
|
||||
temp_llm_index_dir: Path,
|
||||
settings: SettingsWrapper,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
_seed_bloated_index(temp_llm_index_dir)
|
||||
settings.LLM_INDEX_COMPACTION_LOCK_TIMEOUT = 0.3
|
||||
|
||||
lock = _reader_lock()
|
||||
with lock.read_lock(), caplog.at_level(logging.INFO):
|
||||
indexing.llm_index_compact() # must not raise
|
||||
lock.close()
|
||||
|
||||
# Swap was skipped: bloat remains, nothing corrupted, data intact.
|
||||
assert _integrity_ok(temp_llm_index_dir)
|
||||
assert _bloat_ratio(temp_llm_index_dir) > 2
|
||||
assert "Skipping LLM index compaction" in caplog.text
|
||||
|
||||
def test_compaction_runs_when_no_reader_holds_the_lock(
|
||||
self,
|
||||
temp_llm_index_dir: Path,
|
||||
) -> None:
|
||||
_seed_bloated_index(temp_llm_index_dir)
|
||||
assert _bloat_ratio(temp_llm_index_dir) > 2
|
||||
|
||||
indexing.llm_index_compact()
|
||||
|
||||
assert _bloat_ratio(temp_llm_index_dir) == pytest.approx(1.0)
|
||||
assert _integrity_ok(temp_llm_index_dir)
|
||||
|
||||
def test_normal_write_is_not_gated_by_the_compaction_lock(
|
||||
self,
|
||||
temp_llm_index_dir: Path,
|
||||
) -> None:
|
||||
"""A held exclusive lock must not block ordinary writes (WAL handles them)."""
|
||||
_seed_bloated_index(temp_llm_index_dir)
|
||||
done = threading.Event()
|
||||
|
||||
def remove() -> None:
|
||||
indexing.llm_index_remove_document(MagicMock(id=999))
|
||||
done.set()
|
||||
|
||||
holder = _reader_lock()
|
||||
with holder.write_lock():
|
||||
t = threading.Thread(target=remove)
|
||||
t.start()
|
||||
finished = done.wait(timeout=5)
|
||||
t.join(timeout=2)
|
||||
holder.close()
|
||||
assert finished, "a normal write blocked on the compaction lock"
|
||||
|
||||
|
||||
class TestReadStore:
|
||||
def test_closes_connection_on_exit(self, temp_llm_index_dir: Path) -> None:
|
||||
with indexing.read_store() as store:
|
||||
conn = store.client
|
||||
assert conn.execute("SELECT 1").fetchone()[0] == 1
|
||||
with pytest.raises(sqlite3.ProgrammingError):
|
||||
conn.execute("SELECT 1")
|
||||
|
||||
def test_concurrent_readers_do_not_block(self, temp_llm_index_dir: Path) -> None:
|
||||
_seed_bloated_index(temp_llm_index_dir)
|
||||
with indexing.read_store() as a, indexing.read_store() as b:
|
||||
assert a.table_exists()
|
||||
assert b.table_exists()
|
||||
@@ -12,7 +12,7 @@ class TestLazyAiImports:
|
||||
"os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'paperless.settings')\n"
|
||||
"django.setup()\n"
|
||||
"import documents.tasks # noqa: F401\n"
|
||||
"leaked = [m for m in ('lancedb', 'pyarrow', 'llama_index') "
|
||||
"leaked = [m for m in ('lancedb', 'pyarrow', 'llama_index', 'sqlite_vec') "
|
||||
"if m in sys.modules]\n"
|
||||
"assert not leaked, f'AI libraries leaked into the light path: {leaked}'\n"
|
||||
)
|
||||
|
||||
@@ -1,417 +1,606 @@
|
||||
import sqlite3
|
||||
from collections.abc import Generator
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from llama_index.core.schema import NodeRelationship
|
||||
from llama_index.core.schema import RelatedNodeInfo
|
||||
from llama_index.core.schema import TextNode
|
||||
from llama_index.core.vector_stores.types import FilterOperator
|
||||
from llama_index.core.vector_stores.types import MetadataFilter
|
||||
from llama_index.core.vector_stores.types import MetadataFilters
|
||||
from llama_index.core.vector_stores.types import VectorStoreQuery
|
||||
|
||||
from paperless_ai.vector_store import PaperlessLanceVectorStore
|
||||
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
|
||||
from paperless_ai.vector_store import _build_where
|
||||
|
||||
DIM = 8
|
||||
DIM = 16
|
||||
|
||||
|
||||
def _node(node_id: str, document_id: str, text: str, vec: float) -> TextNode:
|
||||
node = TextNode(id_=node_id, text=text, metadata={"document_id": document_id})
|
||||
node.set_content(text)
|
||||
node.embedding = [vec] * DIM
|
||||
# Use relationships so ref_doc_id resolves correctly (it's a read-only property)
|
||||
node.relationships = {
|
||||
NodeRelationship.SOURCE: RelatedNodeInfo(node_id=document_id),
|
||||
}
|
||||
def make_node(
|
||||
node_id: str,
|
||||
document_id: str,
|
||||
*,
|
||||
modified: str = "2026-06-10T00:00:00",
|
||||
seed: float = 0.0,
|
||||
text: str = "some text",
|
||||
) -> TextNode:
|
||||
node = TextNode(
|
||||
id_=node_id,
|
||||
text=text,
|
||||
metadata={"document_id": document_id, "modified": modified},
|
||||
)
|
||||
node.relationships = {}
|
||||
node.embedding = [seed + i / 100 for i in range(DIM)]
|
||||
return node
|
||||
|
||||
|
||||
class TestPaperlessLanceVectorStoreCrud:
|
||||
@pytest.fixture
|
||||
def store(self, tmp_path: Path) -> PaperlessLanceVectorStore:
|
||||
return PaperlessLanceVectorStore(uri=str(tmp_path / "idx"))
|
||||
@pytest.fixture
|
||||
def store(tmp_path: Path) -> Generator[PaperlessSqliteVecVectorStore, None, None]:
|
||||
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as store:
|
||||
yield store
|
||||
|
||||
def test_add_then_query_returns_node(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
store.add([_node("1-0", "1", "alpha", 0.1), _node("2-0", "2", "beta", 0.9)])
|
||||
|
||||
result = store.query(
|
||||
VectorStoreQuery(query_embedding=[0.1] * DIM, similarity_top_k=1),
|
||||
)
|
||||
def _query(
|
||||
store: PaperlessSqliteVecVectorStore,
|
||||
embedding: list[float],
|
||||
top_k: int = 5,
|
||||
filters=None,
|
||||
):
|
||||
return store.query(
|
||||
VectorStoreQuery(
|
||||
query_embedding=embedding,
|
||||
similarity_top_k=top_k,
|
||||
filters=filters,
|
||||
),
|
||||
)
|
||||
|
||||
assert len(result.nodes) == 1
|
||||
|
||||
def _eq_filter(key: str, value: str):
|
||||
return MetadataFilters(
|
||||
filters=[MetadataFilter(key=key, operator=FilterOperator.EQ, value=value)],
|
||||
)
|
||||
|
||||
|
||||
def _in_filter(document_ids: list[str]):
|
||||
return MetadataFilters(
|
||||
filters=[
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.IN,
|
||||
value=document_ids,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class TestCrud:
|
||||
def test_add_then_query_returns_node(self, store) -> None:
|
||||
node = make_node("n1", "1")
|
||||
assert store.add([node]) == ["n1"]
|
||||
result = _query(store, node.embedding, top_k=1)
|
||||
assert result.ids == ["n1"]
|
||||
assert result.nodes[0].metadata["document_id"] == "1"
|
||||
# cosine distance of the identical vector is 0 -> similarity 1
|
||||
assert result.similarities[0] == pytest.approx(1.0)
|
||||
|
||||
def test_query_empty_table_returns_empty_no_raise(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
result = store.query(
|
||||
VectorStoreQuery(query_embedding=[0.1] * DIM, similarity_top_k=5),
|
||||
)
|
||||
assert result.nodes == []
|
||||
assert result.ids == []
|
||||
def test_query_empty_store_returns_empty_no_raise(self, store) -> None:
|
||||
result = _query(store, [0.0] * DIM)
|
||||
assert result.ids == [] and result.nodes == [] and result.similarities == []
|
||||
|
||||
def test_delete_removes_all_chunks_of_document(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
store.add([_node("1-0", "1", "a", 0.1), _node("1-1", "1", "b", 0.2)])
|
||||
store.add([_node("2-0", "2", "c", 0.9)])
|
||||
def test_add_empty_list_is_noop(self, store) -> None:
|
||||
assert store.add([]) == []
|
||||
assert not store.table_exists()
|
||||
|
||||
def test_delete_removes_all_chunks_of_document(self, store) -> None:
|
||||
store.add([make_node("a1", "1"), make_node("a2", "1"), make_node("b1", "2")])
|
||||
store.delete("1")
|
||||
result = _query(store, [0.0] * DIM, top_k=10)
|
||||
assert result.ids == ["b1"]
|
||||
|
||||
assert store.client.open_table("documents").count_rows() == 1
|
||||
def test_query_with_in_filter_scopes_results(self, store) -> None:
|
||||
store.add(
|
||||
[
|
||||
make_node("a1", "1", seed=0.0),
|
||||
make_node("b1", "2", seed=1.0),
|
||||
make_node("c1", "3", seed=2.0),
|
||||
],
|
||||
)
|
||||
result = _query(store, [0.0] * DIM, top_k=10, filters=_in_filter(["2", "3"]))
|
||||
assert sorted(result.ids) == ["b1", "c1"]
|
||||
|
||||
def test_query_with_in_filter_scopes_results(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
store.add([_node("1-0", "1", "a", 0.1), _node("2-0", "2", "b", 0.1)])
|
||||
def test_query_respects_top_k_with_filter(self, store) -> None:
|
||||
# k semantics: global top-k even with IN filters (document_id is a
|
||||
# metadata column, not a partition key -- see design doc).
|
||||
store.add(
|
||||
[make_node(f"n{i}", str(i % 4), seed=float(i)) for i in range(12)],
|
||||
)
|
||||
result = _query(
|
||||
store,
|
||||
[0.0] * DIM,
|
||||
top_k=3,
|
||||
filters=_in_filter(["0", "1", "2", "3"]),
|
||||
)
|
||||
assert len(result.ids) == 3
|
||||
assert result.similarities == sorted(result.similarities, reverse=True)
|
||||
|
||||
result = store.query(
|
||||
VectorStoreQuery(
|
||||
query_embedding=[0.1] * DIM,
|
||||
similarity_top_k=5,
|
||||
filters=MetadataFilters(
|
||||
filters=[
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.IN,
|
||||
value=["2"],
|
||||
),
|
||||
],
|
||||
def test_get_nodes_filter_and_empty_paths(self, store) -> None:
|
||||
assert store.get_nodes(filters=_in_filter(["1"])) == [] # no table yet
|
||||
store.add([make_node("a1", "1"), make_node("b1", "2")])
|
||||
nodes = store.get_nodes(filters=_in_filter(["1"]))
|
||||
assert [n.node_id for n in nodes] == ["a1"]
|
||||
assert nodes[0].embedding is not None
|
||||
assert store.get_nodes(filters=_in_filter(["999"])) == []
|
||||
|
||||
def test_query_with_eq_filter_scopes_results(self, store) -> None:
|
||||
store.add(
|
||||
[
|
||||
make_node("a1", "1", seed=0.0),
|
||||
make_node("b1", "2", seed=1.0),
|
||||
make_node("c1", "3", seed=2.0),
|
||||
],
|
||||
)
|
||||
result = _query(
|
||||
store,
|
||||
[0.0] * DIM,
|
||||
top_k=10,
|
||||
filters=_eq_filter("document_id", "2"),
|
||||
)
|
||||
assert result.ids == ["b1"]
|
||||
|
||||
def test_get_nodes_node_ids_not_implemented(self, store) -> None:
|
||||
with pytest.raises(NotImplementedError):
|
||||
store.get_nodes(node_ids=["x"])
|
||||
|
||||
def test_fresh_instance_sees_existing_table(self, store, tmp_path: Path) -> None:
|
||||
store.add([make_node("a1", "1")])
|
||||
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reopened:
|
||||
assert reopened.table_exists()
|
||||
assert reopened.vector_dim() == DIM
|
||||
assert _query(reopened, [0.0] * DIM, top_k=1).ids == ["a1"]
|
||||
|
||||
def test_table_exists_and_drop(self, store) -> None:
|
||||
assert not store.table_exists()
|
||||
store.add([make_node("a1", "1")])
|
||||
assert store.table_exists()
|
||||
store.drop_table()
|
||||
assert not store.table_exists()
|
||||
assert store.vector_dim() is None
|
||||
|
||||
|
||||
class TestBuildWhere:
|
||||
def test_fails_closed_when_no_filter_is_translatable(self) -> None:
|
||||
# A nested MetadataFilters is not a MetadataFilter, so it is skipped.
|
||||
# With no translatable clauses, the function must fail closed rather
|
||||
# than emit "()" (invalid SQL) and never widen document access.
|
||||
nested = MetadataFilters(
|
||||
filters=[
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.EQ,
|
||||
value="1",
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
where, params = _build_where(MetadataFilters(filters=[nested]))
|
||||
assert where == "1 = 0"
|
||||
assert params == []
|
||||
|
||||
assert [n.metadata["document_id"] for n in result.nodes] == ["2"]
|
||||
def test_query_with_untranslatable_filter_returns_no_rows(self, store) -> None:
|
||||
store.add([make_node("a1", "1"), make_node("b1", "2")])
|
||||
nested = MetadataFilters(
|
||||
filters=[
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.EQ,
|
||||
value="1",
|
||||
),
|
||||
],
|
||||
)
|
||||
filters = MetadataFilters(filters=[nested])
|
||||
# Must not raise (no "WHERE ()") and must return nothing (fail closed).
|
||||
assert _query(store, [0.0] * DIM, top_k=5, filters=filters).ids == []
|
||||
assert store.get_nodes(filters=filters) == []
|
||||
|
||||
def test_get_nodes_filter_returns_empty_cleanly(
|
||||
|
||||
class TestUpsert:
|
||||
def test_upsert_replaces_and_prunes_stale_chunks(self, store) -> None:
|
||||
store.add(
|
||||
[make_node("d1c1", "1"), make_node("d1c2", "1"), make_node("d2c1", "2")],
|
||||
)
|
||||
store.upsert_document("1", [make_node("d1new", "1")])
|
||||
result = _query(store, [0.0] * DIM, top_k=10)
|
||||
assert sorted(result.ids) == ["d1new", "d2c1"]
|
||||
|
||||
def test_upsert_creates_table_when_missing(self, store) -> None:
|
||||
store.upsert_document("1", [make_node("a1", "1")])
|
||||
assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"]
|
||||
|
||||
def test_upsert_empty_nodes_removes_document(self, store) -> None:
|
||||
store.add([make_node("a1", "1"), make_node("b1", "2")])
|
||||
store.upsert_document("1", [])
|
||||
assert _query(store, [0.0] * DIM, top_k=10).ids == ["b1"]
|
||||
|
||||
def test_upsert_is_atomic_for_concurrent_readers(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
store,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
store.add([_node("1-0", "1", "a", 0.1)])
|
||||
nodes = store.get_nodes(
|
||||
filters=MetadataFilters(
|
||||
filters=[
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.IN,
|
||||
value=["999"],
|
||||
),
|
||||
],
|
||||
),
|
||||
)
|
||||
assert nodes == []
|
||||
"""A second connection must never observe document 1 half-replaced."""
|
||||
store.add([make_node("a1", "1"), make_node("a2", "1")])
|
||||
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reader:
|
||||
store.upsert_document("1", [make_node("a3", "1")])
|
||||
ids = [n.node_id for n in reader.get_nodes(filters=_in_filter(["1"]))]
|
||||
assert ids == ["a3"]
|
||||
|
||||
def test_get_nodes_returns_empty_when_no_table(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
result = store.get_nodes(
|
||||
filters=MetadataFilters(
|
||||
filters=[
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.IN,
|
||||
value=["1"],
|
||||
),
|
||||
],
|
||||
),
|
||||
)
|
||||
assert result == []
|
||||
|
||||
def test_fresh_instance_filters_existing_table(
|
||||
class TestMetadataCoercion:
|
||||
def test_none_metadata_values_become_empty_strings(self, store) -> None:
|
||||
node = make_node("a1", "1")
|
||||
node.metadata["modified"] = None
|
||||
store.add([node]) # must not raise (vec0 rejects NULL metadata)
|
||||
assert store.get_modified_times() == {"1": ""}
|
||||
|
||||
|
||||
class TestModelNameTracking:
|
||||
def test_stored_model_name_none_without_table(self, tmp_path: Path) -> None:
|
||||
with PaperlessSqliteVecVectorStore(
|
||||
uri=str(tmp_path),
|
||||
embed_model_name="model-a",
|
||||
) as store:
|
||||
assert store.stored_model_name() is None
|
||||
|
||||
def test_model_name_stored_after_add_and_persists(self, tmp_path: Path) -> None:
|
||||
with PaperlessSqliteVecVectorStore(
|
||||
uri=str(tmp_path),
|
||||
embed_model_name="model-a",
|
||||
) as store:
|
||||
store.add([make_node("a1", "1")])
|
||||
assert store.stored_model_name() == "model-a"
|
||||
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reopened:
|
||||
assert reopened.stored_model_name() == "model-a"
|
||||
|
||||
def test_config_mismatch_semantics(self, tmp_path: Path) -> None:
|
||||
with PaperlessSqliteVecVectorStore(
|
||||
uri=str(tmp_path),
|
||||
embed_model_name="model-a",
|
||||
) as store:
|
||||
assert not store.config_mismatch("anything") # no table yet
|
||||
store.add([make_node("a1", "1")])
|
||||
assert not store.config_mismatch("model-a")
|
||||
assert store.config_mismatch("model-b")
|
||||
|
||||
def test_config_mismatch_false_when_table_predates_tracking(
|
||||
self,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
uri = str(tmp_path / "idx")
|
||||
PaperlessLanceVectorStore(uri=uri).add(
|
||||
[_node("1-0", "1", "a", 0.1), _node("2-0", "2", "b", 0.1)],
|
||||
)
|
||||
|
||||
reopened = PaperlessLanceVectorStore(uri=uri)
|
||||
result = reopened.query(
|
||||
VectorStoreQuery(
|
||||
query_embedding=[0.1] * DIM,
|
||||
similarity_top_k=5,
|
||||
filters=MetadataFilters(
|
||||
filters=[
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.IN,
|
||||
value=["1"],
|
||||
),
|
||||
],
|
||||
),
|
||||
),
|
||||
)
|
||||
assert [n.metadata["document_id"] for n in result.nodes] == ["1"]
|
||||
|
||||
def test_table_exists_and_drop(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
assert store.table_exists() is False
|
||||
store.add([_node("1-0", "1", "a", 0.1)])
|
||||
assert store.table_exists() is True
|
||||
assert store.vector_dim() == DIM
|
||||
store.drop_table()
|
||||
assert store.table_exists() is False
|
||||
|
||||
def test_build_where_or_condition(self) -> None:
|
||||
from llama_index.core.vector_stores.types import FilterCondition
|
||||
|
||||
from paperless_ai.vector_store import _build_where
|
||||
|
||||
where = _build_where(
|
||||
MetadataFilters(
|
||||
filters=[
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.EQ,
|
||||
value="1",
|
||||
),
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.EQ,
|
||||
value="2",
|
||||
),
|
||||
],
|
||||
condition=FilterCondition.OR,
|
||||
),
|
||||
)
|
||||
assert where == "document_id = '1' OR document_id = '2'"
|
||||
|
||||
|
||||
class TestPaperlessLanceVectorStoreUpsert:
|
||||
@pytest.fixture
|
||||
def store(self, tmp_path: Path) -> PaperlessLanceVectorStore:
|
||||
s = PaperlessLanceVectorStore(uri=str(tmp_path / "idx"))
|
||||
s.add(
|
||||
[
|
||||
_node("1-0", "1", "old0", 0.1),
|
||||
_node("1-1", "1", "old1", 0.2),
|
||||
_node("1-2", "1", "old2", 0.3),
|
||||
_node("2-0", "2", "keep", 0.9),
|
||||
],
|
||||
)
|
||||
return s
|
||||
|
||||
def test_upsert_prunes_stale_chunks_and_keeps_others(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
store.upsert_document(
|
||||
"1",
|
||||
[_node("1-0", "1", "new0", 0.1), _node("1-1", "1", "new1", 0.2)],
|
||||
)
|
||||
|
||||
table = store.client.open_table("documents")
|
||||
doc1 = sorted(
|
||||
r["id"] for r in table.search().where("document_id = '1'").to_list()
|
||||
)
|
||||
assert doc1 == ["1-0", "1-1"] # 1-2 pruned
|
||||
assert table.count_rows() == 3 # 2 new doc1 + 1 doc2
|
||||
|
||||
def test_upsert_is_single_commit(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
table = store.client.open_table("documents")
|
||||
before = table.version
|
||||
store.upsert_document("1", [_node("1-0", "1", "new0", 0.1)])
|
||||
assert store.client.open_table("documents").version == before + 1
|
||||
|
||||
def test_upsert_empty_nodes_removes_document(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
store.upsert_document("1", [])
|
||||
|
||||
table = store.client.open_table("documents")
|
||||
remaining = sorted(r["document_id"] for r in table.search().to_list())
|
||||
assert "1" not in remaining
|
||||
assert "2" in remaining
|
||||
|
||||
|
||||
class TestPaperlessLanceVectorStoreMaintenance:
|
||||
@pytest.fixture
|
||||
def store(self, tmp_path: Path) -> PaperlessLanceVectorStore:
|
||||
return PaperlessLanceVectorStore(uri=str(tmp_path / "idx"))
|
||||
|
||||
def test_maybe_create_ann_index_noop_below_threshold(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
store.add([_node("1-0", "1", "a", 0.1)])
|
||||
# Threshold far above row count -> no index attempted, no error.
|
||||
store.maybe_create_ann_index(min_rows=1000)
|
||||
# Still queryable.
|
||||
result = store.query(
|
||||
VectorStoreQuery(query_embedding=[0.1] * DIM, similarity_top_k=1),
|
||||
)
|
||||
assert len(result.nodes) == 1
|
||||
|
||||
def test_maybe_create_ann_index_non_divisible_dim_falls_back(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
# DIM=8 is not divisible by the PQ default sub-vectors; must not raise
|
||||
# and must leave the table queryable (IVF_FLAT fallback or skipped).
|
||||
for i in range(40):
|
||||
store.add([_node(f"1-{i}", "1", f"t{i}", float(i))])
|
||||
store.maybe_create_ann_index(min_rows=10)
|
||||
result = store.query(
|
||||
VectorStoreQuery(query_embedding=[1.0] * DIM, similarity_top_k=3),
|
||||
)
|
||||
assert len(result.nodes) == 3
|
||||
|
||||
def test_compact_reduces_to_single_version(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
for i in range(5):
|
||||
store.add([_node(f"1-{i}", "1", f"t{i}", float(i))])
|
||||
assert len(store.client.open_table("documents").list_versions()) > 1
|
||||
store.compact(retention_seconds=0)
|
||||
assert len(store.client.open_table("documents").list_versions()) == 1
|
||||
|
||||
def test_upsert_after_optimize_with_scalar_index(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
store.add(
|
||||
[
|
||||
_node("1-0", "1", "old0", 0.1),
|
||||
_node("1-1", "1", "old1", 0.2),
|
||||
_node("1-2", "1", "old2", 0.3),
|
||||
_node("2-0", "2", "keep", 0.9),
|
||||
],
|
||||
)
|
||||
store.ensure_document_id_scalar_index()
|
||||
store.compact(retention_seconds=0)
|
||||
|
||||
store.upsert_document("1", [_node("1-0", "1", "new0", 0.1)])
|
||||
|
||||
table = store.client.open_table("documents")
|
||||
doc1 = sorted(
|
||||
r["id"] for r in table.search().where("document_id = '1'").to_list()
|
||||
)
|
||||
assert doc1 == ["1-0"]
|
||||
assert table.count_rows() == 2
|
||||
|
||||
def test_ensure_scalar_index_is_idempotent(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
store.add([_node("1-0", "1", "text", 0.5)])
|
||||
store.ensure_document_id_scalar_index()
|
||||
# Second call must not raise and must not replace the existing index.
|
||||
store.ensure_document_id_scalar_index()
|
||||
assert store._has_index_on("document_id")
|
||||
|
||||
def test_ensure_scalar_index_noop_on_empty_store(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
store.ensure_document_id_scalar_index() # no table yet — must not raise
|
||||
|
||||
|
||||
class TestConfigMismatch:
|
||||
@pytest.fixture
|
||||
def uri(self, tmp_path: Path) -> str:
|
||||
return str(tmp_path / "idx")
|
||||
|
||||
def test_stored_model_name_returns_none_when_no_table(self, uri: str) -> None:
|
||||
store = PaperlessLanceVectorStore(uri=uri)
|
||||
assert store.stored_model_name() is None
|
||||
|
||||
def test_model_name_stored_in_schema_after_add(self, uri: str) -> None:
|
||||
store = PaperlessLanceVectorStore(uri=uri, embed_model_name="all-MiniLM-L6-v2")
|
||||
store.add([_node("1-0", "1", "text", 0.1)])
|
||||
assert store.stored_model_name() == "all-MiniLM-L6-v2"
|
||||
|
||||
def test_model_name_stored_in_schema_after_upsert(self, uri: str) -> None:
|
||||
store = PaperlessLanceVectorStore(uri=uri, embed_model_name="nomic-embed")
|
||||
store.upsert_document("1", [_node("1-0", "1", "text", 0.1)])
|
||||
assert store.stored_model_name() == "nomic-embed"
|
||||
|
||||
def test_model_name_persists_after_reopen(self, uri: str) -> None:
|
||||
PaperlessLanceVectorStore(uri=uri, embed_model_name="all-MiniLM-L6-v2").add(
|
||||
[_node("1-0", "1", "text", 0.1)],
|
||||
)
|
||||
reopened = PaperlessLanceVectorStore(uri=uri)
|
||||
assert reopened.stored_model_name() == "all-MiniLM-L6-v2"
|
||||
|
||||
def test_config_mismatch_returns_false_when_no_table(self, uri: str) -> None:
|
||||
store = PaperlessLanceVectorStore(uri=uri)
|
||||
assert store.config_mismatch("any-model") is False
|
||||
|
||||
def test_config_mismatch_returns_false_when_model_matches(self, uri: str) -> None:
|
||||
store = PaperlessLanceVectorStore(uri=uri, embed_model_name="all-MiniLM-L6-v2")
|
||||
store.add([_node("1-0", "1", "text", 0.1)])
|
||||
assert store.config_mismatch("all-MiniLM-L6-v2") is False
|
||||
|
||||
def test_config_mismatch_returns_true_when_model_differs(self, uri: str) -> None:
|
||||
store = PaperlessLanceVectorStore(uri=uri, embed_model_name="old-model")
|
||||
store.add([_node("1-0", "1", "text", 0.1)])
|
||||
assert store.config_mismatch("new-model") is True
|
||||
|
||||
def test_config_mismatch_returns_false_when_no_metadata_stored(
|
||||
self,
|
||||
uri: str,
|
||||
) -> None:
|
||||
# Tables created before model-name tracking was added have no schema metadata.
|
||||
# Conservative default: assume compatible rather than force a rebuild.
|
||||
store = PaperlessLanceVectorStore(uri=uri)
|
||||
store.add([_node("1-0", "1", "text", 0.1)])
|
||||
assert store.config_mismatch("any-model") is False
|
||||
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as store: # no model name
|
||||
store.add([make_node("a1", "1")])
|
||||
assert not store.config_mismatch("model-a")
|
||||
|
||||
|
||||
class TestGetModifiedTimes:
|
||||
@pytest.fixture
|
||||
def store(self, tmp_path: Path) -> PaperlessLanceVectorStore:
|
||||
return PaperlessLanceVectorStore(uri=str(tmp_path / "idx"))
|
||||
|
||||
def _node_with_modified(
|
||||
self,
|
||||
node_id: str,
|
||||
doc_id: str,
|
||||
modified: str,
|
||||
) -> TextNode:
|
||||
node = TextNode(
|
||||
id_=node_id,
|
||||
text="text",
|
||||
metadata={"document_id": doc_id, "modified": modified},
|
||||
)
|
||||
node.embedding = [0.1] * DIM
|
||||
node.relationships = {
|
||||
NodeRelationship.SOURCE: RelatedNodeInfo(node_id=doc_id),
|
||||
}
|
||||
return node
|
||||
|
||||
def test_empty_store_returns_empty_dict(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
def test_empty_store_returns_empty_dict(self, store) -> None:
|
||||
assert store.get_modified_times() == {}
|
||||
|
||||
def test_returns_one_entry_per_document(
|
||||
self,
|
||||
store: PaperlessLanceVectorStore,
|
||||
) -> None:
|
||||
def test_returns_one_entry_per_document(self, store) -> None:
|
||||
store.add(
|
||||
[
|
||||
self._node_with_modified("1-0", "1", "2024-01-01T00:00:00"),
|
||||
self._node_with_modified("1-1", "1", "2024-01-01T00:00:00"),
|
||||
self._node_with_modified("2-0", "2", "2024-06-01T00:00:00"),
|
||||
make_node("a1", "1", modified="2026-01-01T00:00:00"),
|
||||
make_node("a2", "1", modified="2026-01-01T00:00:00"),
|
||||
make_node("b1", "2", modified="2026-02-02T00:00:00"),
|
||||
],
|
||||
)
|
||||
result = store.get_modified_times()
|
||||
assert result == {
|
||||
"1": "2024-01-01T00:00:00",
|
||||
"2": "2024-06-01T00:00:00",
|
||||
assert store.get_modified_times() == {
|
||||
"1": "2026-01-01T00:00:00",
|
||||
"2": "2026-02-02T00:00:00",
|
||||
}
|
||||
|
||||
|
||||
class TestCompact:
|
||||
def _bloat_ratio(self, store) -> float:
|
||||
live = store.client.execute(
|
||||
"SELECT count(*) FROM documents",
|
||||
).fetchone()[0]
|
||||
# vec0 0.1.9 does not accumulate deleted rows in the _rowids shadow
|
||||
# table, so we track cumulative inserts in index_meta instead.
|
||||
row = store.client.execute(
|
||||
"SELECT value FROM index_meta WHERE key = 'total_inserts'",
|
||||
).fetchone()
|
||||
total = int(row["value"]) if row else live
|
||||
return total / max(live, 1)
|
||||
|
||||
def _churn(self, store, cycles: int) -> None:
|
||||
for i in range(cycles):
|
||||
store.upsert_document(
|
||||
"1",
|
||||
[make_node(f"gen{i}-{j}", "1", seed=float(j)) for j in range(20)],
|
||||
)
|
||||
|
||||
def test_compact_noop_below_threshold(self, store) -> None:
|
||||
store.add([make_node("a1", "1")])
|
||||
store.compact()
|
||||
assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"]
|
||||
|
||||
def test_force_compact_preserves_rows_and_metadata(self, store) -> None:
|
||||
store.add([make_node("a1", "1"), make_node("b1", "2", seed=3.0)])
|
||||
self._churn(store, 5)
|
||||
before = {
|
||||
n.node_id: n.metadata
|
||||
for n in store.get_nodes(filters=_in_filter(["1", "2"]))
|
||||
}
|
||||
store.compact(force=True)
|
||||
after = {
|
||||
n.node_id: n.metadata
|
||||
for n in store.get_nodes(filters=_in_filter(["1", "2"]))
|
||||
}
|
||||
assert after == before
|
||||
assert self._bloat_ratio(store) == pytest.approx(1.0)
|
||||
# store remains fully usable after the rebuild; use a seed far from all
|
||||
# existing nodes (gen4-0..gen4-19 have seeds 0..19) so cosine KNN is
|
||||
# unambiguous at top_k=1.
|
||||
store.upsert_document("3", [make_node("c1", "3", seed=100.0)])
|
||||
assert "c1" in _query(store, [100.0] * DIM, top_k=1).ids
|
||||
|
||||
def test_auto_compact_triggers_on_churn(self, store) -> None:
|
||||
store.add([make_node(f"s{j}", "1", seed=float(j)) for j in range(20)])
|
||||
self._churn(store, 5)
|
||||
assert self._bloat_ratio(store) > 2
|
||||
store.compact()
|
||||
assert self._bloat_ratio(store) == pytest.approx(1.0)
|
||||
|
||||
def test_compact_on_missing_table_is_noop(self, store) -> None:
|
||||
store.compact()
|
||||
store.compact(force=True)
|
||||
|
||||
def test_failed_compact_removes_temp_wal_and_shm(
|
||||
self,
|
||||
store,
|
||||
tmp_path: Path,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
"""A compact() that raises mid-rebuild must leave no .compact* files.
|
||||
|
||||
Normally the sole connection's close() checkpoints the temp WAL away,
|
||||
but a concurrent reader keeps -wal/-shm alive, so the cleanup must
|
||||
unlink them explicitly (as the structural-migration path does).
|
||||
"""
|
||||
store.add([make_node("a1", "1")])
|
||||
compact_path = str(tmp_path / DB_FILENAME) + ".compact"
|
||||
held: list[sqlite3.Connection] = []
|
||||
|
||||
def boom(conn: sqlite3.Connection, dim: int) -> None:
|
||||
# Hold an extra connection so close() of the rebuild connection is
|
||||
# not the last one -> the temp -wal/-shm survive the checkpoint.
|
||||
extra = sqlite3.connect(compact_path)
|
||||
extra.execute("SELECT 1").fetchall()
|
||||
held.append(extra)
|
||||
raise RuntimeError("boom")
|
||||
|
||||
monkeypatch.setattr(
|
||||
PaperlessSqliteVecVectorStore,
|
||||
"_create_vec_table",
|
||||
staticmethod(boom),
|
||||
)
|
||||
try:
|
||||
with pytest.raises(RuntimeError):
|
||||
store.compact(force=True)
|
||||
assert sorted(p.name for p in tmp_path.glob("*.compact*")) == []
|
||||
finally:
|
||||
for c in held:
|
||||
c.close()
|
||||
|
||||
|
||||
class TestDbFile:
|
||||
def test_single_db_file_in_index_dir(self, store, tmp_path: Path) -> None:
|
||||
store.add([make_node("a1", "1")])
|
||||
assert (tmp_path / DB_FILENAME).exists()
|
||||
|
||||
def test_wal_mode_enabled(self, store) -> None:
|
||||
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
|
||||
|
||||
def test_stop_at_reembed_boundary(self, store) -> None:
|
||||
# Registry: structural v2, re-embed v3, structural v4.
|
||||
# Only v2 should apply; the re-embed boundary must stop execution
|
||||
# before v4 runs, and the stored version must stay at 2.
|
||||
store.add([make_node("a1", "1"), make_node("b1", "2")])
|
||||
|
||||
def copy_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("COMMIT")
|
||||
|
||||
migrations = [
|
||||
Migration(
|
||||
from_version=1,
|
||||
to_version=2,
|
||||
kind="structural",
|
||||
description="v2 structural",
|
||||
apply=copy_apply,
|
||||
),
|
||||
Migration(
|
||||
from_version=2,
|
||||
to_version=3,
|
||||
kind="re-embed",
|
||||
description="v3 re-embed boundary",
|
||||
),
|
||||
Migration(
|
||||
from_version=3,
|
||||
to_version=4,
|
||||
kind="structural",
|
||||
description="v4 structural - must not run",
|
||||
apply=copy_apply,
|
||||
),
|
||||
]
|
||||
MIGRATIONS.extend(migrations)
|
||||
try:
|
||||
from paperless_ai import vector_store as vs_mod
|
||||
|
||||
original = vs_mod.SCHEMA_VERSION
|
||||
vs_mod.SCHEMA_VERSION = 4
|
||||
result = store.check_and_run_migrations()
|
||||
finally:
|
||||
for m in migrations:
|
||||
MIGRATIONS.remove(m)
|
||||
vs_mod.SCHEMA_VERSION = original
|
||||
|
||||
assert result is True
|
||||
assert self._schema_version(store) == 2
|
||||
+447
-176
@@ -1,15 +1,25 @@
|
||||
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 types import TracebackType
|
||||
from typing import Any
|
||||
from typing import Literal
|
||||
|
||||
import lancedb
|
||||
import pyarrow as pa
|
||||
import sqlite_vec
|
||||
from llama_index.core.bridge.pydantic import PrivateAttr
|
||||
from llama_index.core.schema import BaseNode
|
||||
from llama_index.core.vector_stores.types import BasePydanticVectorStore
|
||||
from llama_index.core.vector_stores.types import FilterCondition
|
||||
from llama_index.core.vector_stores.types import FilterOperator
|
||||
from llama_index.core.vector_stores.types import MetadataFilter
|
||||
from llama_index.core.vector_stores.types import MetadataFilters
|
||||
from llama_index.core.vector_stores.types import VectorStoreQuery
|
||||
from llama_index.core.vector_stores.types import VectorStoreQueryResult
|
||||
@@ -18,46 +28,118 @@ from llama_index.core.vector_stores.utils import node_to_metadata_dict
|
||||
|
||||
logger = logging.getLogger("paperless_ai.vector_store")
|
||||
|
||||
DB_FILENAME = "llmindex.db"
|
||||
DEFAULT_TABLE_NAME = "documents"
|
||||
|
||||
# Below this many chunks, LanceDB's exact (brute-force) search is sufficient and
|
||||
# faster than building an ANN index (per LanceDB guidance, ~100K vectors).
|
||||
ANN_INDEX_MIN_ROWS = 100_000
|
||||
# IVF_PQ default; num_sub_vectors must evenly divide the embedding dimension.
|
||||
ANN_PQ_SUB_VECTORS = 96
|
||||
# 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
|
||||
# a rebuild copies the live rows into a fresh table.
|
||||
COMPACT_BLOAT_RATIO = 2.0
|
||||
|
||||
# Filterable vec0 metadata columns. _build_where() only ever receives filter
|
||||
# keys we construct ourselves, but allowlisting keeps SQL identifiers safe by
|
||||
# construction.
|
||||
_FILTER_COLUMNS = frozenset({"document_id", "modified"})
|
||||
|
||||
|
||||
def _escape(value: str) -> str:
|
||||
return str(value).replace("'", "''")
|
||||
@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,
|
||||
)
|
||||
|
||||
|
||||
def _build_where(filters: MetadataFilters | None) -> str | None:
|
||||
"""Translate the EQ / IN filters we use into a Lance SQL predicate on the
|
||||
top-level ``document_id`` column."""
|
||||
# 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)
|
||||
|
||||
|
||||
def _unpack(blob: bytes) -> list[float]:
|
||||
return list(struct.unpack(f"{len(blob) // 4}f", blob))
|
||||
|
||||
|
||||
def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]:
|
||||
"""Translate the EQ / IN filters we use into a parameterized SQL clause
|
||||
on vec0 metadata columns. Returns ("", []) when there is nothing to filter.
|
||||
"""
|
||||
if filters is None or not filters.filters:
|
||||
return None
|
||||
return "", []
|
||||
clauses: list[str] = []
|
||||
params: list[str] = []
|
||||
for f in filters.filters:
|
||||
# filters.filters is Union[MetadataFilter, ExactMatchFilter, MetadataFilters];
|
||||
# we only build MetadataFilter entries, so skip anything else at runtime.
|
||||
if not isinstance(f, MetadataFilter):
|
||||
continue
|
||||
if f.key not in _FILTER_COLUMNS: # pragma: no cover - we build the keys
|
||||
raise NotImplementedError(f"Unsupported filter column: {f.key}")
|
||||
if f.operator == FilterOperator.IN:
|
||||
vals = ",".join(f"'{_escape(v)}'" for v in f.value)
|
||||
clauses.append(f"{f.key} IN ({vals})")
|
||||
values = [str(v) for v in f.value] # type: ignore[union-attr] # value is list when operator is IN
|
||||
if not values: # pragma: no cover
|
||||
clauses.append("1 = 0")
|
||||
continue
|
||||
placeholders = ",".join("?" for _ in values)
|
||||
clauses.append(f"{f.key} IN ({placeholders})")
|
||||
params.extend(values)
|
||||
elif f.operator == FilterOperator.EQ:
|
||||
clauses.append(f"{f.key} = '{_escape(f.value)}'")
|
||||
clauses.append(f"{f.key} = ?")
|
||||
params.append(str(f.value))
|
||||
else: # pragma: no cover - we only ever build EQ/IN filters
|
||||
raise NotImplementedError(f"Unsupported filter operator: {f.operator}")
|
||||
if not clauses:
|
||||
# Filters were requested but none could be translated. Fail closed
|
||||
# rather than emit "()" (invalid SQL): filters scope document access,
|
||||
# so an empty translation must match no rows, never widen the scope.
|
||||
return "1 = 0", []
|
||||
joiner = " OR " if filters.condition == FilterCondition.OR else " AND "
|
||||
return joiner.join(clauses)
|
||||
return "(" + joiner.join(clauses) + ")", params
|
||||
|
||||
|
||||
class PaperlessLanceVectorStore(BasePydanticVectorStore):
|
||||
"""A llama-index vector store backed directly by a LanceDB table.
|
||||
class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
||||
"""A llama-index vector store backed by a sqlite-vec vec0 table.
|
||||
|
||||
Stores one row per node with the node id, its document id (both as the
|
||||
``ref_doc_id`` delete key ``doc_id`` and a top-level filter column
|
||||
``document_id``), the embedding, and the serialised node (text + metadata)
|
||||
as JSON. ``stores_text`` lets llama-index run off this store alone, with no
|
||||
Stores one row per node: the node id (TEXT primary key), its document id
|
||||
(metadata column, used for EQ/IN filtering and per-document delete), the
|
||||
document's modified timestamp, the embedding (float32, cosine metric), and
|
||||
the serialized node (text + metadata) as JSON in an auxiliary column.
|
||||
``stores_text`` lets llama-index run off this store alone, with no
|
||||
separate docstore or index store.
|
||||
|
||||
Everything lives in one SQLite database file (``DB_FILENAME``) inside the
|
||||
directory given as ``uri`` (kept as a directory for compatibility with the
|
||||
previous LanceDB layout). WAL mode allows readers in other processes to
|
||||
proceed while the (FileLock-serialized) writer holds a transaction.
|
||||
|
||||
Implemented surface of ``BasePydanticVectorStore``
|
||||
---------------------------------------------------
|
||||
Only the methods actively used by this codebase are implemented.
|
||||
@@ -70,58 +152,117 @@ class PaperlessLanceVectorStore(BasePydanticVectorStore):
|
||||
flat_metadata: bool = False
|
||||
|
||||
_uri: str = PrivateAttr()
|
||||
_table_name: str = PrivateAttr()
|
||||
_embed_model_name: str | None = PrivateAttr()
|
||||
_conn: Any = PrivateAttr()
|
||||
_table: Any = PrivateAttr()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
uri: str,
|
||||
table_name: str = DEFAULT_TABLE_NAME,
|
||||
embed_model_name: str | None = None,
|
||||
) -> None:
|
||||
super().__init__(stores_text=True, flat_metadata=False)
|
||||
self._uri = uri
|
||||
self._table_name = table_name
|
||||
self._embed_model_name = embed_model_name
|
||||
self._conn = lancedb.connect(uri)
|
||||
existing = self._conn.list_tables().tables
|
||||
self._table = (
|
||||
self._conn.open_table(table_name) if table_name in existing else None
|
||||
self._conn = self._open_connection(str(Path(uri) / DB_FILENAME))
|
||||
|
||||
@staticmethod
|
||||
def _open_connection(db_path: str) -> sqlite3.Connection:
|
||||
conn = sqlite3.connect(
|
||||
db_path,
|
||||
timeout=30,
|
||||
isolation_level=None, # autocommit; explicit transactions below
|
||||
)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.enable_load_extension(True) # noqa: FBT003
|
||||
sqlite_vec.load(conn)
|
||||
conn.enable_load_extension(False) # noqa: FBT003
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("PRAGMA synchronous=NORMAL")
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS index_meta (key TEXT PRIMARY KEY, value TEXT)",
|
||||
)
|
||||
return conn
|
||||
|
||||
@property
|
||||
def client(self) -> Any:
|
||||
return self._conn
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the underlying SQLite connection (idempotent)."""
|
||||
self._conn.close()
|
||||
|
||||
def __enter__(self) -> "PaperlessSqliteVecVectorStore":
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
# Deterministically release the connection (and its WAL/SHM handles) so
|
||||
# it is never left open across a compaction/migration file swap.
|
||||
self.close()
|
||||
|
||||
@contextmanager
|
||||
def _transaction(self) -> Iterator[None]:
|
||||
self._conn.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
yield
|
||||
except BaseException: # pragma: no cover
|
||||
self._conn.execute("ROLLBACK")
|
||||
raise
|
||||
else:
|
||||
self._conn.execute("COMMIT")
|
||||
|
||||
def _meta_get(self, key: str) -> str | None:
|
||||
row = self._conn.execute(
|
||||
"SELECT value FROM index_meta WHERE key = ?",
|
||||
(key,),
|
||||
).fetchone()
|
||||
return row["value"] if row else None
|
||||
|
||||
@staticmethod
|
||||
def _meta_set_on(conn: sqlite3.Connection, key: str, value: str) -> None:
|
||||
conn.execute(
|
||||
"INSERT INTO index_meta (key, value) VALUES (?, ?) "
|
||||
"ON CONFLICT(key) DO UPDATE SET value = excluded.value",
|
||||
(key, value),
|
||||
)
|
||||
|
||||
def _meta_set(self, key: str, value: str) -> None:
|
||||
self._meta_set_on(self._conn, key, value)
|
||||
|
||||
def table_exists(self) -> bool:
|
||||
return self._table is not None
|
||||
return (
|
||||
self._conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?",
|
||||
(DEFAULT_TABLE_NAME,),
|
||||
).fetchone()
|
||||
is not None
|
||||
)
|
||||
|
||||
def vector_dim(self) -> int | None:
|
||||
if self._table is None:
|
||||
if not self.table_exists():
|
||||
return None
|
||||
return self._table.schema.field("vector").type.list_size
|
||||
value = self._meta_get("dim")
|
||||
return int(value) if value else None
|
||||
|
||||
def drop_table(self) -> None:
|
||||
if self.table_exists():
|
||||
self._conn.drop_table(self._table_name)
|
||||
self._table = None
|
||||
self._conn.execute("DROP TABLE IF EXISTS " + DEFAULT_TABLE_NAME)
|
||||
self._conn.execute("DELETE FROM index_meta")
|
||||
|
||||
def stored_model_name(self) -> str | None:
|
||||
"""Return the embedding model name stored in table schema metadata, or None."""
|
||||
if self._table is None:
|
||||
"""Return the embedding model name recorded at table creation, or None."""
|
||||
if not self.table_exists():
|
||||
return None
|
||||
meta = self._table.schema.metadata or {}
|
||||
value = meta.get(b"embed_model")
|
||||
return value.decode() if value else None
|
||||
return self._meta_get("embed_model")
|
||||
|
||||
def config_mismatch(self, model_name: str) -> bool:
|
||||
"""True when the stored model name differs from ``model_name``.
|
||||
|
||||
Returns False when no table exists or when the table predates model-name
|
||||
tracking (schema has no metadata) — conservative default avoids spurious
|
||||
rebuilds on upgrade.
|
||||
Returns False when no table exists or when the table predates
|
||||
model-name tracking — conservative default avoids spurious rebuilds.
|
||||
"""
|
||||
stored = self.stored_model_name()
|
||||
if stored is None:
|
||||
@@ -129,97 +270,115 @@ class PaperlessLanceVectorStore(BasePydanticVectorStore):
|
||||
return stored != model_name
|
||||
|
||||
@staticmethod
|
||||
def _schema(dim: int, model_name: str | None = None) -> pa.Schema:
|
||||
meta = {b"embed_model": model_name.encode()} if model_name else None
|
||||
return pa.schema(
|
||||
[
|
||||
pa.field("id", pa.string()),
|
||||
pa.field("doc_id", pa.string()),
|
||||
pa.field("document_id", pa.string()),
|
||||
pa.field("modified", pa.string()),
|
||||
pa.field("vector", pa.list_(pa.float32(), dim)),
|
||||
pa.field("node_content", pa.string()),
|
||||
],
|
||||
metadata=meta,
|
||||
def _create_vec_table(conn: sqlite3.Connection, dim: int) -> None:
|
||||
# document_id is deliberately a metadata column, NOT a partition key:
|
||||
# partition keys change KNN `k` to per-partition semantics under IN
|
||||
# filters (asg017/sqlite-vec#142); metadata columns give a correct
|
||||
# global top-k.
|
||||
conn.execute( # nosemgrep: python.sqlalchemy.security.sqlalchemy-execute-raw-query.sqlalchemy-execute-raw-query
|
||||
"CREATE VIRTUAL TABLE "
|
||||
+ DEFAULT_TABLE_NAME
|
||||
+ " USING vec0("
|
||||
+ "id TEXT PRIMARY KEY,"
|
||||
+ " document_id TEXT,"
|
||||
+ " modified TEXT,"
|
||||
+ " +node_content TEXT,"
|
||||
+ " embedding float["
|
||||
+ str(int(dim))
|
||||
+ "] distance_metric=cosine"
|
||||
+ ")",
|
||||
)
|
||||
|
||||
def _row(self, node: BaseNode) -> dict[str, Any]:
|
||||
def _create_table(self, dim: int) -> None:
|
||||
self._create_vec_table(self._conn, dim)
|
||||
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)
|
||||
|
||||
def _ensure_table(self, dim: int) -> None:
|
||||
if not self.table_exists():
|
||||
self._create_table(dim)
|
||||
|
||||
def _row(self, node: BaseNode) -> tuple[str, str, str, str, bytes]:
|
||||
meta = node_to_metadata_dict(
|
||||
node,
|
||||
remove_text=False,
|
||||
flat_metadata=self.flat_metadata,
|
||||
)
|
||||
return {
|
||||
"id": node.node_id,
|
||||
"doc_id": node.ref_doc_id,
|
||||
"document_id": str(node.metadata.get("document_id")),
|
||||
"modified": str(node.metadata.get("modified", "")),
|
||||
"vector": node.get_embedding(),
|
||||
"node_content": json.dumps(meta),
|
||||
}
|
||||
|
||||
def _ensure_table(self, rows: list[dict[str, Any]], dim: int) -> bool:
|
||||
"""Create the table from ``rows`` if it does not exist yet.
|
||||
|
||||
Returns True if the table was just created (caller can skip the
|
||||
separate add/merge step), False if the table already existed.
|
||||
"""
|
||||
if self._table is not None:
|
||||
return False
|
||||
self._table = self._conn.create_table(
|
||||
self._table_name,
|
||||
rows,
|
||||
schema=self._schema(dim, self._embed_model_name),
|
||||
# vec0 metadata columns reject NULL (asg017/sqlite-vec#141): coerce
|
||||
# every value to a string, with "" as the absent sentinel.
|
||||
document_id = node.ref_doc_id or node.metadata.get("document_id")
|
||||
return (
|
||||
node.node_id,
|
||||
str(document_id or ""),
|
||||
str(node.metadata.get("modified") or ""),
|
||||
json.dumps(meta),
|
||||
_pack(node.get_embedding()),
|
||||
)
|
||||
return True
|
||||
|
||||
_INSERT = (
|
||||
"INSERT INTO "
|
||||
+ DEFAULT_TABLE_NAME
|
||||
+ " (id, document_id, modified, node_content, embedding) VALUES (?, ?, ?, ?, ?)"
|
||||
)
|
||||
|
||||
def _increment_total_inserts(self, count: int) -> None:
|
||||
"""Increment the cumulative insert counter stored in index_meta.
|
||||
|
||||
This counter never decreases (DELETEs do not decrement it) and is
|
||||
used by compact() to estimate the bloat ratio: when total_inserts /
|
||||
live_rows exceeds COMPACT_BLOAT_RATIO the table has accumulated
|
||||
enough deleted-but-not-freed rows to warrant a rebuild.
|
||||
"""
|
||||
current = int(self._meta_get("total_inserts") or "0")
|
||||
self._meta_set("total_inserts", str(current + count))
|
||||
|
||||
def add(self, nodes: Sequence[BaseNode], **add_kwargs: Any) -> list[str]:
|
||||
if not nodes:
|
||||
return []
|
||||
rows = [self._row(node) for node in nodes]
|
||||
dim = len(nodes[0].get_embedding())
|
||||
if not self._ensure_table(rows, dim):
|
||||
self._table.add(rows)
|
||||
with self._transaction():
|
||||
self._ensure_table(len(nodes[0].get_embedding()))
|
||||
self._conn.executemany(self._INSERT, rows)
|
||||
self._increment_total_inserts(len(rows))
|
||||
return [node.node_id for node in nodes]
|
||||
|
||||
def upsert_document(self, document_id: str, nodes: list[BaseNode]) -> list[str]:
|
||||
"""Atomically replace all stored chunks of ``document_id`` with ``nodes``.
|
||||
|
||||
A single ``merge_insert`` commit: matching node ids are updated, new ids
|
||||
inserted, and any existing rows for this document that are not in the new
|
||||
set are deleted (``when_not_matched_by_source_delete``). This prunes stale
|
||||
trailing chunks when an edit reduces a document's chunk count, with no
|
||||
transient empty state for concurrent lock-free readers.
|
||||
One transaction deletes the document's existing rows and inserts the
|
||||
new set (vec0's INSERT OR REPLACE is broken upstream, #259, so
|
||||
delete+insert it is). WAL readers in other processes see either the
|
||||
old or the new chunk set, never a partial state.
|
||||
"""
|
||||
if not nodes:
|
||||
# No indexable content: remove any existing chunks for this document.
|
||||
if self._table is not None:
|
||||
self._table.delete(f"document_id = '{_escape(document_id)}'")
|
||||
return []
|
||||
rows = [self._row(node) for node in nodes]
|
||||
dim = len(nodes[0].get_embedding())
|
||||
if self._ensure_table(rows, dim):
|
||||
return [node.node_id for node in nodes]
|
||||
(
|
||||
self._table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.when_not_matched_by_source_delete(
|
||||
f"document_id = '{_escape(document_id)}'",
|
||||
)
|
||||
.execute(rows)
|
||||
)
|
||||
with self._transaction():
|
||||
if nodes:
|
||||
self._ensure_table(len(nodes[0].get_embedding()))
|
||||
if self.table_exists():
|
||||
self._conn.execute(
|
||||
"DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE document_id = ?",
|
||||
(str(document_id),),
|
||||
)
|
||||
if rows:
|
||||
self._conn.executemany(self._INSERT, rows)
|
||||
self._increment_total_inserts(len(rows))
|
||||
return [node.node_id for node in nodes]
|
||||
|
||||
def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
|
||||
if self._table is not None:
|
||||
self._table.delete(f"doc_id = '{_escape(ref_doc_id)}'")
|
||||
if self.table_exists():
|
||||
with self._transaction():
|
||||
self._conn.execute(
|
||||
"DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE document_id = ?",
|
||||
(str(ref_doc_id),),
|
||||
)
|
||||
|
||||
def _rows_to_nodes(self, rows: list[dict[str, Any]]) -> list[BaseNode]:
|
||||
def _rows_to_nodes(self, rows: list[sqlite3.Row]) -> list[BaseNode]:
|
||||
nodes: list[BaseNode] = []
|
||||
for row in rows:
|
||||
node = metadata_dict_to_node(json.loads(row["node_content"]))
|
||||
node.embedding = list(row["vector"])
|
||||
node.embedding = _unpack(row["embedding"])
|
||||
nodes.append(node)
|
||||
return nodes
|
||||
|
||||
@@ -232,102 +391,214 @@ class PaperlessLanceVectorStore(BasePydanticVectorStore):
|
||||
if node_ids is not None: # pragma: no cover
|
||||
# node_ids lookup is not implemented; see class docstring.
|
||||
raise NotImplementedError(
|
||||
"PaperlessLanceVectorStore does not support node_ids lookup",
|
||||
"PaperlessSqliteVecVectorStore does not support node_ids lookup",
|
||||
)
|
||||
if self._table is None:
|
||||
if not self.table_exists():
|
||||
return []
|
||||
where = _build_where(filters)
|
||||
query = self._table.search()
|
||||
where, params = _build_where(filters)
|
||||
sql = "SELECT node_content, embedding FROM " + DEFAULT_TABLE_NAME
|
||||
if where:
|
||||
query = query.where(where)
|
||||
return self._rows_to_nodes(query.to_list())
|
||||
sql += " WHERE " + where
|
||||
return self._rows_to_nodes(self._conn.execute(sql, params).fetchall())
|
||||
|
||||
def query(
|
||||
self,
|
||||
query: VectorStoreQuery,
|
||||
**kwargs: Any,
|
||||
) -> VectorStoreQueryResult:
|
||||
if self._table is None:
|
||||
if not self.table_exists():
|
||||
return VectorStoreQueryResult(nodes=[], similarities=[], ids=[])
|
||||
if query.query_embedding is None: # pragma: no cover
|
||||
return VectorStoreQueryResult(nodes=[], similarities=[], ids=[])
|
||||
top_k = query.similarity_top_k if query.similarity_top_k is not None else 10
|
||||
search = self._table.search(query.query_embedding).limit(top_k)
|
||||
where = _build_where(query.filters)
|
||||
where, params = _build_where(query.filters)
|
||||
sql = (
|
||||
"SELECT id, node_content, embedding, distance FROM "
|
||||
+ DEFAULT_TABLE_NAME
|
||||
+ " WHERE embedding MATCH ? AND k = ?"
|
||||
)
|
||||
if where:
|
||||
search = search.where(where)
|
||||
rows = search.to_list()
|
||||
sql += " AND " + where
|
||||
rows = self._conn.execute(
|
||||
sql,
|
||||
[_pack(query.query_embedding), top_k, *params],
|
||||
).fetchall()
|
||||
# vec0 returns rows distance-sorted ascending; slice defensively in
|
||||
# case future schema changes alter k semantics (e.g. partition keys
|
||||
# return k rows per partition).
|
||||
rows = rows[:top_k]
|
||||
nodes = self._rows_to_nodes(rows)
|
||||
# LanceDB returns an L2 distance (smaller = closer); map to a descending similarity.
|
||||
sims = [1.0 / (1.0 + float(row["_distance"])) for row in rows]
|
||||
# Cosine distance in [0, 2]; map to a descending similarity.
|
||||
# vec0 returns None distance when the query embedding is the zero vector
|
||||
# (no meaningful cosine angle); treat that as maximum distance (1.0) so
|
||||
# the row is included but ranked last.
|
||||
sims = [
|
||||
1.0 - float(row["distance"] if row["distance"] is not None else 1.0)
|
||||
for row in rows
|
||||
]
|
||||
ids = [row["id"] for row in rows]
|
||||
return VectorStoreQueryResult(nodes=nodes, similarities=sims, ids=ids)
|
||||
|
||||
def _has_index_on(self, column: str) -> bool:
|
||||
return any(column in idx.columns for idx in self._table.list_indices())
|
||||
|
||||
def maybe_create_ann_index(self, min_rows: int = ANN_INDEX_MIN_ROWS) -> None:
|
||||
"""Best-effort: build an IVF index once the table is large enough.
|
||||
|
||||
IVF_PQ is used when ``num_sub_vectors`` divides the embedding dimension,
|
||||
otherwise IVF_FLAT (no divisor constraint). Any failure is logged and
|
||||
leaves the table on exact search, which is always correct.
|
||||
"""
|
||||
if self._table is None:
|
||||
return
|
||||
rows = self._table.count_rows()
|
||||
if rows < min_rows or self._has_index_on("vector"):
|
||||
return
|
||||
num_partitions = max(1, rows // 4096)
|
||||
# Embedding dim from the schema's fixed-size list column.
|
||||
dim = self._table.schema.field("vector").type.list_size
|
||||
try:
|
||||
if dim % ANN_PQ_SUB_VECTORS == 0: # pragma: no cover
|
||||
self._table.create_index(
|
||||
metric="l2",
|
||||
num_partitions=num_partitions,
|
||||
num_sub_vectors=ANN_PQ_SUB_VECTORS,
|
||||
index_type="IVF_PQ",
|
||||
)
|
||||
else:
|
||||
self._table.create_index(
|
||||
metric="l2",
|
||||
num_partitions=num_partitions,
|
||||
index_type="IVF_FLAT",
|
||||
)
|
||||
except Exception as e: # pragma: no cover - depends on data/dim
|
||||
logger.warning("Skipping ANN index creation: %s", e)
|
||||
|
||||
def get_modified_times(self) -> dict[str, str]:
|
||||
"""Return {document_id: stored_modified_isoformat} for all indexed documents.
|
||||
|
||||
One representative chunk per document is fetched; all chunks share the
|
||||
same ``modified`` value so the first one seen is sufficient.
|
||||
All chunks of a document share the same ``modified`` value, so the
|
||||
first row seen per document is sufficient.
|
||||
"""
|
||||
if self._table is None:
|
||||
if not self.table_exists():
|
||||
return {}
|
||||
result: dict[str, str] = {}
|
||||
for row in self._table.search().select(["document_id", "modified"]).to_list():
|
||||
for row in self._conn.execute(
|
||||
"SELECT document_id, modified FROM " + DEFAULT_TABLE_NAME,
|
||||
):
|
||||
doc_id = str(row["document_id"])
|
||||
if doc_id not in result:
|
||||
result[doc_id] = str(row["modified"] or "")
|
||||
return result
|
||||
|
||||
def ensure_document_id_scalar_index(self) -> None:
|
||||
"""Create a scalar index on the filter column (never on the merge key
|
||||
``id`` — see https://github.com/lancedb/lancedb/issues/3177).
|
||||
No-op if the index already exists."""
|
||||
if self._table is None:
|
||||
def compact(self, *, force: bool = False) -> None:
|
||||
"""Rebuild the database file to reclaim space left behind by DELETEs.
|
||||
|
||||
vec0 DELETE only invalidates rows; the vector data stays in the file
|
||||
forever (asg017/sqlite-vec#54), and per-document re-indexing is a
|
||||
delete+insert. The cumulative insert counter in ``index_meta`` tracks
|
||||
total rows ever written; when that exceeds ``COMPACT_BLOAT_RATIO`` x
|
||||
the live row count (or when forced), live rows are copied into a fresh
|
||||
database file and swapped in via ``os.replace``.
|
||||
|
||||
Note: ``ALTER TABLE ... RENAME TO`` on vec0 virtual tables does NOT
|
||||
rename the shadow tables (sqlite-vec upstream limitation), so
|
||||
an in-place rename-based rebuild is not safe. The file-swap approach
|
||||
is the maintainer-endorsed workaround (asg017/sqlite-vec#205).
|
||||
"""
|
||||
if not self.table_exists():
|
||||
return
|
||||
if self._has_index_on("document_id"):
|
||||
live = self._conn.execute(
|
||||
"SELECT count(*) FROM " + DEFAULT_TABLE_NAME,
|
||||
).fetchone()[0]
|
||||
total = int(self._meta_get("total_inserts") or str(live))
|
||||
if not force and total <= max(live, 1) * COMPACT_BLOAT_RATIO:
|
||||
return
|
||||
dim = self.vector_dim()
|
||||
if dim is None: # pragma: no cover - dim is written at creation
|
||||
logger.warning("Skipping compact: no stored vector dimension")
|
||||
return
|
||||
logger.info(
|
||||
"Compacting LLM index (%d live rows, %d cumulative inserts)",
|
||||
live,
|
||||
total,
|
||||
)
|
||||
db_path = str(Path(self._uri) / DB_FILENAME)
|
||||
compact_path = db_path + ".compact"
|
||||
|
||||
# Copy all live rows into a fresh database file.
|
||||
new_conn = self._open_connection(compact_path)
|
||||
try:
|
||||
self._table.create_scalar_index("document_id")
|
||||
except Exception as e: # pragma: no cover
|
||||
logger.warning("Skipping document_id scalar index: %s", e)
|
||||
self._create_vec_table(new_conn, dim)
|
||||
self._meta_set_on(new_conn, "dim", str(dim))
|
||||
for key in ("embed_model", "schema_version"):
|
||||
value = self._meta_get(key)
|
||||
if value is not None:
|
||||
self._meta_set_on(new_conn, key, value)
|
||||
rows = self._conn.execute(
|
||||
"SELECT id, document_id, modified, node_content, embedding "
|
||||
"FROM " + DEFAULT_TABLE_NAME,
|
||||
).fetchall()
|
||||
new_conn.execute("BEGIN IMMEDIATE")
|
||||
new_conn.executemany(
|
||||
self._INSERT,
|
||||
[
|
||||
(
|
||||
r["id"],
|
||||
r["document_id"],
|
||||
r["modified"],
|
||||
r["node_content"],
|
||||
bytes(r["embedding"]),
|
||||
)
|
||||
for r in rows
|
||||
],
|
||||
)
|
||||
# Reset the cumulative counter: after compact, total_inserts == live.
|
||||
self._meta_set_on(new_conn, "total_inserts", str(live))
|
||||
new_conn.execute("COMMIT")
|
||||
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._swap_in_compact(compact_path, db_path)
|
||||
|
||||
def compact(self, retention_seconds: int) -> None:
|
||||
"""Compact fragments and prune old MVCC versions in one call."""
|
||||
if self._table is None:
|
||||
return
|
||||
from datetime import timedelta
|
||||
def _swap_in_compact(self, compact_path: str, db_path: str) -> None:
|
||||
"""Atomically replace the live database with the compacted copy."""
|
||||
self._conn.close()
|
||||
for suffix in ["-wal", "-shm"]:
|
||||
stale = Path(compact_path + suffix)
|
||||
if stale.exists(): # pragma: no cover
|
||||
stale.unlink()
|
||||
Path(compact_path).replace(db_path)
|
||||
self._conn = self._open_connection(db_path)
|
||||
|
||||
self._table.optimize(cleanup_older_than=timedelta(seconds=retention_seconds))
|
||||
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)
|
||||
|
||||
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)
|
||||
self._meta_set_on(new_conn, "schema_version", str(migration.to_version))
|
||||
except BaseException: # pragma: no cover
|
||||
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._swap_in_compact(compact_path, db_path)
|
||||
Reference in new issue
Block a user