mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-08-01 08:32:18 +00:00
Compare commits
9
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0c1e10914f | ||
|
|
a6c3139d2a | ||
|
|
d56f54001f | ||
|
|
25365f9cbb | ||
|
|
2c8bc3d5cf | ||
|
|
831ee05b7c | ||
|
|
1261cdf5ce | ||
|
|
1e3c9a762f | ||
|
|
601dd5d05d |
@@ -39,7 +39,6 @@ from documents.search._tokenizer import ascii_fold
|
|||||||
from documents.search._tokenizer import autocomplete_tokens
|
from documents.search._tokenizer import autocomplete_tokens
|
||||||
from documents.search._tokenizer import register_tokenizers
|
from documents.search._tokenizer import register_tokenizers
|
||||||
from documents.utils import IterWrapper
|
from documents.utils import IterWrapper
|
||||||
from documents.utils import QuerySetStream
|
|
||||||
from documents.utils import identity
|
from documents.utils import identity
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -1036,15 +1035,14 @@ _EMPTY_VIEWER_GRANT: Final[ViewerGrant] = ViewerGrant(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
class _DocumentViewerStream(QuerySetStream["Document"]):
|
class _DocumentViewerStream:
|
||||||
"""Yield document permission data while batch-loading grants.
|
"""Yield document permission data while batch-loading grants.
|
||||||
|
|
||||||
Viewer permissions are fetched in batches (see
|
Viewer permissions are fetched in batches (see
|
||||||
``_bulk_get_viewer_permissions``), but documents are yielded individually so a
|
``_bulk_get_viewer_permissions``), but documents are yielded individually so a
|
||||||
progress bar wrapped around this stream advances per document rather than
|
progress bar wrapped around this stream advances per document rather than
|
||||||
jumping a whole chunk at a time. ``__len__`` (inherited from
|
jumping a whole chunk at a time. ``__len__`` lets the progress helper still
|
||||||
``QuerySetStream``) lets the progress helper still discover the total (it
|
discover the total (it inspects ``QuerySet``/``Sized``).
|
||||||
inspects ``QuerySet``/``Sized``).
|
|
||||||
|
|
||||||
The viewer and group ids travel with each document in the yielded pair
|
The viewer and group ids travel with each document in the yielded pair
|
||||||
rather than through a separate mutable attribute, so the pairing survives
|
rather than through a separate mutable attribute, so the pairing survives
|
||||||
@@ -1053,11 +1051,18 @@ class _DocumentViewerStream(QuerySetStream["Document"]):
|
|||||||
generator in lock-step.
|
generator in lock-step.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
def __init__(self, documents: QuerySet[Document], *, chunk_size: int) -> None:
|
||||||
|
self._documents = documents
|
||||||
|
self._chunk_size = chunk_size
|
||||||
|
|
||||||
|
def __len__(self) -> int:
|
||||||
|
return self._documents.count()
|
||||||
|
|
||||||
def __iter__(self) -> Iterator[tuple[Document, ViewerGrant]]:
|
def __iter__(self) -> Iterator[tuple[Document, ViewerGrant]]:
|
||||||
# iterator(chunk_size=…) streams from a server-side cursor instead of
|
# iterator(chunk_size=…) streams from a server-side cursor instead of
|
||||||
# materialising the whole queryset in memory; since Django 4.1 it still
|
# materialising the whole queryset in memory; since Django 4.1 it still
|
||||||
# honours prefetch_related, running the prefetches one batch at a time.
|
# honours prefetch_related, running the prefetches one batch at a time.
|
||||||
documents = self._queryset.iterator(chunk_size=self._chunk_size)
|
documents = self._documents.iterator(chunk_size=self._chunk_size)
|
||||||
for chunk in chunked(documents, self._chunk_size):
|
for chunk in chunked(documents, self._chunk_size):
|
||||||
grants_by_pk = _bulk_get_viewer_permissions([doc.pk for doc in chunk])
|
grants_by_pk = _bulk_get_viewer_permissions([doc.pk for doc in chunk])
|
||||||
for doc in chunk:
|
for doc in chunk:
|
||||||
|
|||||||
@@ -1,35 +0,0 @@
|
|||||||
import pytest_mock
|
|
||||||
|
|
||||||
from documents.utils import QuerySetStream
|
|
||||||
|
|
||||||
|
|
||||||
class TestQuerySetStream:
|
|
||||||
def test_len_and_iter_delegate_to_streaming_queryset_methods(
|
|
||||||
self,
|
|
||||||
mocker: pytest_mock.MockerFixture,
|
|
||||||
) -> None:
|
|
||||||
"""
|
|
||||||
GIVEN:
|
|
||||||
- A mock queryset
|
|
||||||
WHEN:
|
|
||||||
- A QuerySetStream wrapping it is measured and iterated
|
|
||||||
THEN:
|
|
||||||
- len() uses count() (not a materializing len()), and iteration
|
|
||||||
uses .iterator(chunk_size=...) (not plain iteration, which
|
|
||||||
would materialize the whole queryset, plus any prefetch
|
|
||||||
caches, into Django's own result cache at once)
|
|
||||||
"""
|
|
||||||
mock_queryset = mocker.MagicMock()
|
|
||||||
mock_queryset.count.return_value = 42
|
|
||||||
mock_queryset.iterator.return_value = iter(["row-1", "row-2"])
|
|
||||||
streamed = QuerySetStream(mock_queryset, chunk_size=1000)
|
|
||||||
|
|
||||||
assert len(streamed) == 42
|
|
||||||
assert list(streamed) == ["row-1", "row-2"]
|
|
||||||
# count.call_count isn't asserted exactly: list()'s own size-hint
|
|
||||||
# optimization calls len(streamed) again internally, on top of the
|
|
||||||
# explicit len() call above -- both legitimately delegate to
|
|
||||||
# count(), so only the delegation itself (not the call count) is
|
|
||||||
# the thing being verified here.
|
|
||||||
mock_queryset.count.assert_called_with()
|
|
||||||
mock_queryset.iterator.assert_called_once_with(chunk_size=1000)
|
|
||||||
@@ -3,24 +3,16 @@ import logging
|
|||||||
import shutil
|
import shutil
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
from collections.abc import Iterable
|
from collections.abc import Iterable
|
||||||
from collections.abc import Iterator
|
|
||||||
from os import utime
|
from os import utime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from subprocess import CompletedProcess
|
from subprocess import CompletedProcess
|
||||||
from subprocess import run
|
from subprocess import run
|
||||||
from typing import TYPE_CHECKING
|
|
||||||
from typing import Generic
|
|
||||||
from typing import TypeVar
|
from typing import TypeVar
|
||||||
|
|
||||||
from django.conf import settings
|
from django.conf import settings
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from django.db.models import Model
|
|
||||||
from django.db.models import QuerySet
|
|
||||||
|
|
||||||
_T = TypeVar("_T")
|
_T = TypeVar("_T")
|
||||||
_M = TypeVar("_M", bound="Model")
|
|
||||||
|
|
||||||
# A function that wraps an iterable — typically used to inject a progress bar.
|
# A function that wraps an iterable — typically used to inject a progress bar.
|
||||||
IterWrapper = Callable[[Iterable[_T]], Iterable[_T]]
|
IterWrapper = Callable[[Iterable[_T]], Iterable[_T]]
|
||||||
@@ -31,40 +23,6 @@ def identity(iterable: Iterable[_T]) -> Iterable[_T]:
|
|||||||
return iterable
|
return iterable
|
||||||
|
|
||||||
|
|
||||||
class QuerySetStream(Generic[_M]):
|
|
||||||
"""Stream a QuerySet via .iterator(chunk_size=...) instead of
|
|
||||||
materializing it (plus any prefetch caches) all at once, while still
|
|
||||||
supporting len() via count() so a progress bar wrapped around this
|
|
||||||
(e.g. via IterWrapper) shows a real total instead of falling back to
|
|
||||||
indeterminate.
|
|
||||||
|
|
||||||
Plain QuerySet iteration (``for row in queryset:``) is not lazy: Django
|
|
||||||
fetches every matching row in one query and caches the fully-hydrated
|
|
||||||
result in the queryset's own ``_result_cache`` before yielding the
|
|
||||||
first item -- wrapping that in a progress bar or any other iterable
|
|
||||||
adapter doesn't change this, since none of them alter how the
|
|
||||||
underlying queryset produces items. ``.iterator(chunk_size=...)`` is
|
|
||||||
the specific Django API that bypasses ``_result_cache`` and streams
|
|
||||||
from a server-side cursor instead, discarding each chunk once consumed
|
|
||||||
(and, since Django 4.1, still honours ``prefetch_related``, running the
|
|
||||||
prefetches one batch at a time rather than for the whole queryset).
|
|
||||||
|
|
||||||
Subclass to layer additional per-batch work on top (see
|
|
||||||
``documents.search._backend._DocumentViewerStream``) by overriding
|
|
||||||
``__iter__`` -- ``__len__`` and the constructor are inherited for free.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, queryset: "QuerySet[_M]", *, chunk_size: int) -> None:
|
|
||||||
self._queryset = queryset
|
|
||||||
self._chunk_size = chunk_size
|
|
||||||
|
|
||||||
def __len__(self) -> int:
|
|
||||||
return self._queryset.count()
|
|
||||||
|
|
||||||
def __iter__(self) -> Iterator[_M]:
|
|
||||||
return iter(self._queryset.iterator(chunk_size=self._chunk_size))
|
|
||||||
|
|
||||||
|
|
||||||
def _coerce_to_path(
|
def _coerce_to_path(
|
||||||
source: Path | str,
|
source: Path | str,
|
||||||
dest: Path | str,
|
dest: Path | str,
|
||||||
|
|||||||
+111
-28
@@ -1,5 +1,7 @@
|
|||||||
|
import enum
|
||||||
import logging
|
import logging
|
||||||
from collections.abc import Iterable
|
from collections.abc import Iterable
|
||||||
|
from collections.abc import Iterator
|
||||||
from contextlib import contextmanager
|
from contextlib import contextmanager
|
||||||
from datetime import timedelta
|
from datetime import timedelta
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
@@ -14,7 +16,6 @@ from filelock import Timeout
|
|||||||
from documents.models import Document
|
from documents.models import Document
|
||||||
from documents.models import PaperlessTask
|
from documents.models import PaperlessTask
|
||||||
from documents.utils import IterWrapper
|
from documents.utils import IterWrapper
|
||||||
from documents.utils import QuerySetStream
|
|
||||||
from documents.utils import identity
|
from documents.utils import identity
|
||||||
from paperless.config import AIConfig
|
from paperless.config import AIConfig
|
||||||
from paperless_ai.db import db_connection_released
|
from paperless_ai.db import db_connection_released
|
||||||
@@ -23,6 +24,7 @@ from paperless_ai.embedding import get_configured_model_name
|
|||||||
from paperless_ai.embedding import get_embedding_model
|
from paperless_ai.embedding import get_embedding_model
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
|
from django.db.models import QuerySet
|
||||||
from llama_index.core.schema import BaseNode
|
from llama_index.core.schema import BaseNode
|
||||||
|
|
||||||
from paperless_ai.vector_store import PaperlessSqliteVecVectorStore
|
from paperless_ai.vector_store import PaperlessSqliteVecVectorStore
|
||||||
@@ -34,11 +36,35 @@ RAG_NUM_OUTPUT = 512
|
|||||||
RAG_CHUNK_OVERLAP = 200
|
RAG_CHUNK_OVERLAP = 200
|
||||||
|
|
||||||
# update_llm_index(): row count per .iterator() batch when streaming
|
# update_llm_index(): row count per .iterator() batch when streaming
|
||||||
# documents for a rebuild/update via QuerySetStream, matching
|
# documents for a rebuild/update, matching _DocumentViewerStream's chunk
|
||||||
# _DocumentViewerStream's chunk size in documents/search/_backend.py.
|
# size in documents/search/_backend.py.
|
||||||
_INDEX_STREAM_CHUNK_SIZE = 1000
|
_INDEX_STREAM_CHUNK_SIZE = 1000
|
||||||
|
|
||||||
|
|
||||||
|
class _StreamedDocuments:
|
||||||
|
"""A thin QuerySet wrapper that streams via ``.iterator()`` instead of
|
||||||
|
materializing every row (plus its ``content`` and prefetch caches) into
|
||||||
|
memory at once, while still supporting ``len()`` so ``iter_wrapper``'s
|
||||||
|
progress bar shows a real total instead of falling back to indeterminate.
|
||||||
|
Same shape as ``documents/search/_backend.py``'s ``_DocumentViewerStream``,
|
||||||
|
just without that class's extra per-batch permission lookup -- nothing
|
||||||
|
here needs one.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, documents: "QuerySet[Document]") -> None:
|
||||||
|
self._documents = documents
|
||||||
|
|
||||||
|
def __len__(self) -> int:
|
||||||
|
return self._documents.count()
|
||||||
|
|
||||||
|
def __iter__(self) -> Iterator[Document]:
|
||||||
|
# iterator(chunk_size=...) streams from a server-side cursor instead
|
||||||
|
# of materializing the whole queryset in memory; since Django 4.1 it
|
||||||
|
# still honours prefetch_related, running the prefetches one batch
|
||||||
|
# at a time.
|
||||||
|
return iter(self._documents.iterator(chunk_size=_INDEX_STREAM_CHUNK_SIZE))
|
||||||
|
|
||||||
|
|
||||||
def queue_llm_index_update_if_needed(*, rebuild: bool, reason: str) -> bool:
|
def queue_llm_index_update_if_needed(*, rebuild: bool, reason: str) -> bool:
|
||||||
# NOTE: The check-then-enqueue sequence below is non-atomic (TOCTOU): two
|
# NOTE: The check-then-enqueue sequence below is non-atomic (TOCTOU): two
|
||||||
# concurrent workers can both observe no running task and both enqueue a
|
# concurrent workers can both observe no running task and both enqueue a
|
||||||
@@ -192,18 +218,49 @@ def write_store(embed_model_name: str | None = None):
|
|||||||
yield store
|
yield store
|
||||||
|
|
||||||
|
|
||||||
def _check_and_run_migrations(store: "PaperlessSqliteVecVectorStore") -> bool:
|
class MigrationCheckResult(enum.Enum):
|
||||||
"""Run any pending structural migrations, returning True if a pending
|
"""Outcome of _check_and_run_migrations().
|
||||||
re-embed migration needs the caller to force a rebuild -- never
|
|
||||||
triggered automatically here. Safe to call before any write, including
|
CURRENT: no migration was pending, or a pending structural migration
|
||||||
|
was applied successfully -- safe to write.
|
||||||
|
|
||||||
|
REEMBED_REQUIRED: a pending migration needs fresh embeddings, which is
|
||||||
|
never triggered automatically -- the caller must force a rebuild.
|
||||||
|
|
||||||
|
DEFERRED: a migration was pending but could not run because active
|
||||||
|
index readers did not drain within LLM_INDEX_COMPACTION_LOCK_TIMEOUT --
|
||||||
|
the store is still on its old schema. Callers must NOT proceed to
|
||||||
|
write: collapsing this into the same falsy value as CURRENT (as a
|
||||||
|
plain bool return once did) would let a write proceed against an
|
||||||
|
unmigrated schema.
|
||||||
|
"""
|
||||||
|
|
||||||
|
CURRENT = "current"
|
||||||
|
REEMBED_REQUIRED = "reembed_required"
|
||||||
|
DEFERRED = "deferred"
|
||||||
|
|
||||||
|
|
||||||
|
def _check_and_run_migrations(
|
||||||
|
store: "PaperlessSqliteVecVectorStore",
|
||||||
|
) -> MigrationCheckResult:
|
||||||
|
"""Run any pending structural migrations, reporting the outcome as a
|
||||||
|
tri-state result. Safe to call before any write, including
|
||||||
delete()/upsert_document(): has_pending_migration() (see its docstring)
|
delete()/upsert_document(): has_pending_migration() (see its docstring)
|
||||||
keeps this a no-op, with no exclusive access taken, once the store is
|
keeps this a no-op, with no exclusive access taken, once the store is
|
||||||
current.
|
current.
|
||||||
"""
|
"""
|
||||||
if not store.has_pending_migration():
|
if not store.has_pending_migration():
|
||||||
return False
|
return MigrationCheckResult.CURRENT
|
||||||
return bool(
|
result = _with_exclusive_access(
|
||||||
_with_exclusive_access("migration check", store.check_and_run_migrations),
|
"migration check",
|
||||||
|
store.check_and_run_migrations,
|
||||||
|
)
|
||||||
|
if result is None:
|
||||||
|
return MigrationCheckResult.DEFERRED
|
||||||
|
return (
|
||||||
|
MigrationCheckResult.REEMBED_REQUIRED
|
||||||
|
if result
|
||||||
|
else MigrationCheckResult.CURRENT
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -378,12 +435,21 @@ def update_llm_index(
|
|||||||
happens, since a rebuild always covers the whole library regardless.
|
happens, since a rebuild always covers the whole library regardless.
|
||||||
"""
|
"""
|
||||||
with write_store() as store:
|
with write_store() as store:
|
||||||
needs_reembed = _check_and_run_migrations(store)
|
migration_result = _check_and_run_migrations(store)
|
||||||
if needs_reembed:
|
if migration_result is MigrationCheckResult.REEMBED_REQUIRED:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"LLM index migration requires re-embedding; forcing rebuild.",
|
"LLM index migration requires re-embedding; forcing rebuild.",
|
||||||
)
|
)
|
||||||
rebuild = True
|
rebuild = True
|
||||||
|
elif migration_result is MigrationCheckResult.DEFERRED:
|
||||||
|
logger.info(
|
||||||
|
"Skipping LLM index update: migration check deferred while "
|
||||||
|
"index readers are active; will retry next run.",
|
||||||
|
)
|
||||||
|
return (
|
||||||
|
"Skipping LLM index update: migration check deferred; "
|
||||||
|
"will retry next run."
|
||||||
|
)
|
||||||
documents = Document.objects.select_related(
|
documents = Document.objects.select_related(
|
||||||
"correspondent",
|
"correspondent",
|
||||||
"document_type",
|
"document_type",
|
||||||
@@ -399,12 +465,11 @@ def update_llm_index(
|
|||||||
config = AIConfig()
|
config = AIConfig()
|
||||||
model_name = get_configured_model_name(config)
|
model_name = get_configured_model_name(config)
|
||||||
|
|
||||||
if not rebuild and llm_index_exists():
|
if not rebuild:
|
||||||
with read_store() as store:
|
with read_store() as store:
|
||||||
config_mismatch = store.config_mismatch(model_name)
|
if store.table_exists() and store.config_mismatch(model_name):
|
||||||
if config_mismatch:
|
logger.warning("Embedding model changed; forcing LLM index rebuild.")
|
||||||
logger.warning("Embedding model changed; forcing LLM index rebuild.")
|
rebuild = True
|
||||||
rebuild = True
|
|
||||||
|
|
||||||
if no_documents:
|
if no_documents:
|
||||||
logger.warning("No documents found to index.")
|
logger.warning("No documents found to index.")
|
||||||
@@ -416,9 +481,7 @@ def update_llm_index(
|
|||||||
if rebuild or not store.table_exists():
|
if rebuild or not store.table_exists():
|
||||||
logger.info("Rebuilding LLM index.")
|
logger.info("Rebuilding LLM index.")
|
||||||
store.drop_table()
|
store.drop_table()
|
||||||
for document in iter_wrapper(
|
for document in iter_wrapper(_StreamedDocuments(documents)):
|
||||||
QuerySetStream(documents, chunk_size=_INDEX_STREAM_CHUNK_SIZE),
|
|
||||||
):
|
|
||||||
nodes = build_document_node(document, chunk_size=chunk_size)
|
nodes = build_document_node(document, chunk_size=chunk_size)
|
||||||
_embed_nodes(nodes, embed_model)
|
_embed_nodes(nodes, embed_model)
|
||||||
store.add(nodes)
|
store.add(nodes)
|
||||||
@@ -431,9 +494,7 @@ def update_llm_index(
|
|||||||
)
|
)
|
||||||
existing = store.get_modified_times()
|
existing = store.get_modified_times()
|
||||||
changed = 0
|
changed = 0
|
||||||
for document in iter_wrapper(
|
for document in iter_wrapper(_StreamedDocuments(scoped_documents)):
|
||||||
QuerySetStream(scoped_documents, chunk_size=_INDEX_STREAM_CHUNK_SIZE),
|
|
||||||
):
|
|
||||||
doc_id = str(document.id)
|
doc_id = str(document.id)
|
||||||
if existing.get(doc_id) == document.modified.isoformat():
|
if existing.get(doc_id) == document.modified.isoformat():
|
||||||
continue
|
continue
|
||||||
@@ -462,8 +523,8 @@ def llm_index_add_or_update_document(document: Document):
|
|||||||
_embed_nodes(new_nodes, get_embedding_model(config))
|
_embed_nodes(new_nodes, get_embedding_model(config))
|
||||||
|
|
||||||
with write_store(embed_model_name=get_configured_model_name(config)) as store:
|
with write_store(embed_model_name=get_configured_model_name(config)) as store:
|
||||||
needs_reembed = _check_and_run_migrations(store)
|
migration_result = _check_and_run_migrations(store)
|
||||||
if needs_reembed:
|
if migration_result is MigrationCheckResult.REEMBED_REQUIRED:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"Skipping incremental LLM index update for document %s: the "
|
"Skipping incremental LLM index update for document %s: the "
|
||||||
"index requires re-embedding first. Run 'document_llmindex "
|
"index requires re-embedding first. Run 'document_llmindex "
|
||||||
@@ -471,6 +532,14 @@ def llm_index_add_or_update_document(document: Document):
|
|||||||
document.id,
|
document.id,
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
|
if migration_result is MigrationCheckResult.DEFERRED:
|
||||||
|
logger.info(
|
||||||
|
"Skipping incremental LLM index update for document %s: "
|
||||||
|
"migration check deferred while index readers are active; "
|
||||||
|
"will retry on the next write.",
|
||||||
|
document.id,
|
||||||
|
)
|
||||||
|
return
|
||||||
store.upsert_document(str(document.id), new_nodes)
|
store.upsert_document(str(document.id), new_nodes)
|
||||||
|
|
||||||
|
|
||||||
@@ -489,14 +558,19 @@ def llm_index_migrate() -> None:
|
|||||||
if not AIConfig().llm_index_enabled:
|
if not AIConfig().llm_index_enabled:
|
||||||
return
|
return
|
||||||
with write_store() as store:
|
with write_store() as store:
|
||||||
needs_reembed = _check_and_run_migrations(store)
|
migration_result = _check_and_run_migrations(store)
|
||||||
if needs_reembed:
|
if migration_result is MigrationCheckResult.REEMBED_REQUIRED:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"LLM index requires re-embedding, which this automatic migration "
|
"LLM index requires re-embedding, which this automatic migration "
|
||||||
"check will not do on its own -- it can be slow and, for a "
|
"check will not do on its own -- it can be slow and, for a "
|
||||||
"metered embedding backend, cost money. Run "
|
"metered embedding backend, cost money. Run "
|
||||||
"'document_llmindex rebuild' manually when ready.",
|
"'document_llmindex rebuild' manually when ready.",
|
||||||
)
|
)
|
||||||
|
elif migration_result is MigrationCheckResult.DEFERRED:
|
||||||
|
logger.info(
|
||||||
|
"LLM index migration check deferred while index readers are "
|
||||||
|
"active; will retry next run.",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def llm_index_compact() -> None:
|
def llm_index_compact() -> None:
|
||||||
@@ -508,7 +582,8 @@ def llm_index_compact() -> None:
|
|||||||
def llm_index_remove_document(document: Document):
|
def llm_index_remove_document(document: Document):
|
||||||
"""Remove a document's chunks from the LLM index."""
|
"""Remove a document's chunks from the LLM index."""
|
||||||
with write_store() as store:
|
with write_store() as store:
|
||||||
if _check_and_run_migrations(store):
|
migration_result = _check_and_run_migrations(store)
|
||||||
|
if migration_result is MigrationCheckResult.REEMBED_REQUIRED:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"Skipping removal of document %s from the LLM index: the "
|
"Skipping removal of document %s from the LLM index: the "
|
||||||
"index requires re-embedding first. Run 'document_llmindex "
|
"index requires re-embedding first. Run 'document_llmindex "
|
||||||
@@ -516,6 +591,14 @@ def llm_index_remove_document(document: Document):
|
|||||||
document.id,
|
document.id,
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
|
if migration_result is MigrationCheckResult.DEFERRED:
|
||||||
|
logger.info(
|
||||||
|
"Skipping removal of document %s from the LLM index: "
|
||||||
|
"migration check deferred while index readers are active; "
|
||||||
|
"will retry on the next write.",
|
||||||
|
document.id,
|
||||||
|
)
|
||||||
|
return
|
||||||
store.delete(str(document.id))
|
store.delete(str(document.id))
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,118 @@
|
|||||||
|
import sqlite3
|
||||||
|
|
||||||
|
from paperless_ai.migrations import MIGRATIONS
|
||||||
|
from paperless_ai.migrations import Migration
|
||||||
|
from paperless_ai.tables import ChunkRow
|
||||||
|
from paperless_ai.tables import DocumentChunksTable
|
||||||
|
from paperless_ai.tables import DocumentMetaRow
|
||||||
|
from paperless_ai.tables import DocumentMetaTable
|
||||||
|
from paperless_ai.tables import IndexMetaTable
|
||||||
|
from paperless_ai.vector_store import BATCH_SIZE
|
||||||
|
from paperless_ai.vector_store import DEFAULT_TABLE_NAME
|
||||||
|
|
||||||
|
# v1's vec0 shape has never changed since it first shipped and is the ONLY
|
||||||
|
# real upgrade path -- no store has ever existed at any intermediate
|
||||||
|
# version, so this migration goes straight from that shipped shape to the
|
||||||
|
# final v2 target in one pass.
|
||||||
|
_V1_SELECT = (
|
||||||
|
"SELECT id, document_id, modified, node_content, embedding FROM "
|
||||||
|
+ DEFAULT_TABLE_NAME
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _migrate_v1_to_v2(
|
||||||
|
src_conn: sqlite3.Connection,
|
||||||
|
dst_conn: sqlite3.Connection,
|
||||||
|
dim: int,
|
||||||
|
) -> None:
|
||||||
|
"""v1 -> v2: document_id TEXT -> INTEGER, modified moves out of vec0
|
||||||
|
into document_meta, document_chunks added for O(1) per-document delete.
|
||||||
|
|
||||||
|
Freezes its own v2-shaped vec0/document_chunks/document_meta DDL inline,
|
||||||
|
rather than delegating to the gateway "create table" helpers or the
|
||||||
|
store's own vec0-table builder (all of which always reflect the
|
||||||
|
*current* schema): a later schema version changing any of these tables'
|
||||||
|
shape must not silently change what this migration produces for someone
|
||||||
|
upgrading straight from v1.
|
||||||
|
_open_connection() already created document_chunks/document_meta on
|
||||||
|
dst_conn (reflecting current HEAD) as a side effect of opening it for
|
||||||
|
this migration's rebuild -- DROP them first so this migration's own
|
||||||
|
frozen CREATE TABLE isn't a silent no-op against that. Safe here because
|
||||||
|
dst_conn is a freshly opened, empty rebuild file with nothing written
|
||||||
|
yet.
|
||||||
|
"""
|
||||||
|
dst_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 INTEGER,"
|
||||||
|
+ " +node_content TEXT,"
|
||||||
|
+ " embedding float["
|
||||||
|
+ str(int(dim))
|
||||||
|
+ "] distance_metric=cosine"
|
||||||
|
+ ")",
|
||||||
|
)
|
||||||
|
dst_conn.execute("DROP TABLE IF EXISTS document_chunks")
|
||||||
|
dst_conn.execute(
|
||||||
|
"CREATE TABLE document_chunks "
|
||||||
|
"(chunk_id TEXT PRIMARY KEY, document_id INTEGER NOT NULL)",
|
||||||
|
)
|
||||||
|
dst_conn.execute(
|
||||||
|
"CREATE INDEX idx_document_chunks_document_id ON document_chunks (document_id)",
|
||||||
|
)
|
||||||
|
dst_conn.execute("DROP TABLE IF EXISTS document_meta")
|
||||||
|
dst_conn.execute(
|
||||||
|
"CREATE TABLE document_meta "
|
||||||
|
"(document_id INTEGER PRIMARY KEY, modified TEXT NOT NULL)",
|
||||||
|
)
|
||||||
|
|
||||||
|
IndexMetaTable.set_dim(dst_conn, dim)
|
||||||
|
embed_model = IndexMetaTable.get_embed_model(src_conn)
|
||||||
|
if embed_model is not None:
|
||||||
|
IndexMetaTable.set_embed_model(dst_conn, embed_model)
|
||||||
|
|
||||||
|
dst_conn.execute("BEGIN IMMEDIATE")
|
||||||
|
src_cursor = src_conn.execute(_V1_SELECT)
|
||||||
|
live = 0
|
||||||
|
while batch := src_cursor.fetchmany(BATCH_SIZE):
|
||||||
|
vec0_rows = []
|
||||||
|
chunk_rows = []
|
||||||
|
meta_by_document: dict[int, str] = {}
|
||||||
|
for r in batch:
|
||||||
|
document_id = int(r["document_id"])
|
||||||
|
vec0_rows.append(
|
||||||
|
(r["id"], document_id, r["node_content"], bytes(r["embedding"])),
|
||||||
|
)
|
||||||
|
chunk_rows.append(ChunkRow(r["id"], document_id))
|
||||||
|
meta_by_document[document_id] = str(r["modified"] or "")
|
||||||
|
dst_conn.executemany(
|
||||||
|
"INSERT INTO "
|
||||||
|
+ DEFAULT_TABLE_NAME
|
||||||
|
+ " (id, document_id, node_content, embedding) VALUES (?, ?, ?, ?)",
|
||||||
|
vec0_rows,
|
||||||
|
)
|
||||||
|
DocumentChunksTable.insert_many(dst_conn, chunk_rows)
|
||||||
|
DocumentMetaTable.upsert_many(
|
||||||
|
dst_conn,
|
||||||
|
(DocumentMetaRow(doc_id, mod) for doc_id, mod in meta_by_document.items()),
|
||||||
|
)
|
||||||
|
live += len(batch)
|
||||||
|
# This migration only ever copies live rows (like compact()), so the
|
||||||
|
# cumulative counter resets to match -- the new file has no bloat yet.
|
||||||
|
IndexMetaTable.reset_total_inserts(dst_conn, live)
|
||||||
|
dst_conn.execute("COMMIT")
|
||||||
|
|
||||||
|
|
||||||
|
MIGRATIONS.append(
|
||||||
|
Migration(
|
||||||
|
from_version=1,
|
||||||
|
to_version=2,
|
||||||
|
kind="structural",
|
||||||
|
description=(
|
||||||
|
"document_id TEXT -> INTEGER; move modified into document_meta; "
|
||||||
|
"add document_chunks for O(1) per-document delete"
|
||||||
|
),
|
||||||
|
apply=_migrate_v1_to_v2,
|
||||||
|
),
|
||||||
|
)
|
||||||
@@ -0,0 +1,245 @@
|
|||||||
|
"""Thin gateways over the plain relational side tables that sit alongside the
|
||||||
|
vec0 table. Each method takes the sqlite3.Connection to operate on
|
||||||
|
explicitly, rather than owning one -- the store swaps connections during
|
||||||
|
compact()/migration, and migrations always work across two connections
|
||||||
|
(src_conn, dst_conn) at once.
|
||||||
|
|
||||||
|
PRECONDITION: Callers must set conn.row_factory = sqlite3.Row before passing a
|
||||||
|
connection to any of these gateways' read methods. The read methods across all
|
||||||
|
three classes (DocumentChunksTable.chunk_ids_for_document, IndexMetaTable._get,
|
||||||
|
DocumentMetaTable.all_modified_times, DocumentMetaTable.copy_all) use
|
||||||
|
row["column_name"] dictionary-style indexing, which requires sqlite3.Row as the
|
||||||
|
row factory -- without it, sqlite3.Row is not set, rows are returned as plain
|
||||||
|
tuples, and tuple indices must be integers, raising TypeError.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import sqlite3
|
||||||
|
from collections.abc import Iterable
|
||||||
|
from typing import NamedTuple
|
||||||
|
|
||||||
|
|
||||||
|
class ChunkRow(NamedTuple):
|
||||||
|
chunk_id: str
|
||||||
|
document_id: int
|
||||||
|
|
||||||
|
|
||||||
|
class DocumentMetaRow(NamedTuple):
|
||||||
|
document_id: int
|
||||||
|
modified: str
|
||||||
|
|
||||||
|
|
||||||
|
class DocumentChunksTable:
|
||||||
|
"""chunk_id -> document_id, indexed by document_id. Gives O(1)
|
||||||
|
per-document chunk lookup that vec0's own document_id metadata column
|
||||||
|
cannot (see PaperlessSqliteVecVectorStore._delete_chunks_by_document_id).
|
||||||
|
"""
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def create(conn: sqlite3.Connection) -> None:
|
||||||
|
conn.execute(
|
||||||
|
"CREATE TABLE IF NOT EXISTS document_chunks "
|
||||||
|
"(chunk_id TEXT PRIMARY KEY, document_id INTEGER NOT NULL)",
|
||||||
|
)
|
||||||
|
conn.execute(
|
||||||
|
"CREATE INDEX IF NOT EXISTS idx_document_chunks_document_id "
|
||||||
|
"ON document_chunks (document_id)",
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def insert_many(conn: sqlite3.Connection, rows: Iterable[ChunkRow]) -> None:
|
||||||
|
"""rows must already be batch-bounded by the caller (e.g. vec0's own
|
||||||
|
fetchmany() loop) -- this never reads, so it can't itself introduce
|
||||||
|
an unbounded scan, but a whole-table iterable defeats the point."""
|
||||||
|
conn.executemany(
|
||||||
|
"INSERT INTO document_chunks (chunk_id, document_id) VALUES (?, ?)",
|
||||||
|
rows,
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def chunk_ids_for_document(
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
document_id: int,
|
||||||
|
) -> list[str]:
|
||||||
|
return [
|
||||||
|
row["chunk_id"]
|
||||||
|
for row in conn.execute(
|
||||||
|
"SELECT chunk_id FROM document_chunks WHERE document_id = ?",
|
||||||
|
(document_id,),
|
||||||
|
).fetchall()
|
||||||
|
]
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def delete_for_document(conn: sqlite3.Connection, document_id: int) -> None:
|
||||||
|
conn.execute(
|
||||||
|
"DELETE FROM document_chunks WHERE document_id = ?",
|
||||||
|
(document_id,),
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def delete_all(conn: sqlite3.Connection) -> None:
|
||||||
|
conn.execute("DELETE FROM document_chunks")
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def count(conn: sqlite3.Connection) -> int:
|
||||||
|
"""Cheap stand-in for vec0's own row count -- see compact()."""
|
||||||
|
return conn.execute("SELECT count(*) FROM document_chunks").fetchone()[0]
|
||||||
|
|
||||||
|
|
||||||
|
class DocumentMetaTable:
|
||||||
|
"""document_id -> modified, one row per document. Lives outside vec0
|
||||||
|
because vec0 only inlines TEXT metadata up to 12 bytes and `modified`
|
||||||
|
(an ISO timestamp) is always longer.
|
||||||
|
"""
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def create(conn: sqlite3.Connection) -> None:
|
||||||
|
conn.execute(
|
||||||
|
"CREATE TABLE IF NOT EXISTS document_meta "
|
||||||
|
"(document_id INTEGER PRIMARY KEY, modified TEXT NOT NULL)",
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def upsert_many(
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
rows: Iterable[DocumentMetaRow],
|
||||||
|
) -> None:
|
||||||
|
conn.executemany(
|
||||||
|
"INSERT INTO document_meta (document_id, modified) VALUES (?, ?) "
|
||||||
|
"ON CONFLICT(document_id) DO UPDATE SET modified = excluded.modified",
|
||||||
|
rows,
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def delete_for_document(conn: sqlite3.Connection, document_id: int) -> None:
|
||||||
|
conn.execute(
|
||||||
|
"DELETE FROM document_meta WHERE document_id = ?",
|
||||||
|
(document_id,),
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def delete_all(conn: sqlite3.Connection) -> None:
|
||||||
|
conn.execute("DELETE FROM document_meta")
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def copy_all(
|
||||||
|
src_conn: sqlite3.Connection,
|
||||||
|
dst_conn: sqlite3.Connection,
|
||||||
|
batch_size: int,
|
||||||
|
) -> None:
|
||||||
|
"""Stream document_meta from src_conn into dst_conn in bounded
|
||||||
|
batches. The *only* sanctioned way to move this table across
|
||||||
|
connections (compact()/migrations) -- an unbounded fetchall here
|
||||||
|
would defeat the same OOM-avoidance the vec0 row copy already relies
|
||||||
|
on. batch_size has no default: forces the call site to think about
|
||||||
|
it (pass BATCH_SIZE)."""
|
||||||
|
cursor = src_conn.execute(
|
||||||
|
"SELECT document_id, modified FROM document_meta",
|
||||||
|
)
|
||||||
|
while batch := cursor.fetchmany(batch_size):
|
||||||
|
DocumentMetaTable.upsert_many(
|
||||||
|
dst_conn,
|
||||||
|
(DocumentMetaRow(r["document_id"], r["modified"]) for r in batch),
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def all_modified_times(conn: sqlite3.Connection) -> dict[str, str]:
|
||||||
|
"""Full document_id -> modified map, for get_modified_times()'s
|
||||||
|
public API only. One unbounded read by design (existing behavior).
|
||||||
|
Never use this for cross-connection copying; see copy_all()."""
|
||||||
|
return {
|
||||||
|
str(row["document_id"]): str(row["modified"] or "")
|
||||||
|
for row in conn.execute(
|
||||||
|
"SELECT document_id, modified FROM document_meta",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class IndexMetaTable:
|
||||||
|
"""Typed accessors over index_meta's key/value rows -- replaces
|
||||||
|
PaperlessSqliteVecVectorStore._meta_get_on/_meta_set_on, which returned
|
||||||
|
untyped str | None regardless of whether the key held an int (dim,
|
||||||
|
schema_version, total_inserts) or a string (embed_model).
|
||||||
|
"""
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def create(conn: sqlite3.Connection) -> None:
|
||||||
|
conn.execute(
|
||||||
|
"CREATE TABLE IF NOT EXISTS index_meta (key TEXT PRIMARY KEY, value TEXT)",
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _get(conn: sqlite3.Connection, key: str) -> str | None:
|
||||||
|
row = conn.execute(
|
||||||
|
"SELECT value FROM index_meta WHERE key = ?",
|
||||||
|
(key,),
|
||||||
|
).fetchone()
|
||||||
|
return row["value"] if row else None
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _set(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),
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def get_dim(conn: sqlite3.Connection) -> int | None:
|
||||||
|
value = IndexMetaTable._get(conn, "dim")
|
||||||
|
return int(value) if value is not None else None
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def set_dim(conn: sqlite3.Connection, dim: int) -> None:
|
||||||
|
IndexMetaTable._set(conn, "dim", str(dim))
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def get_embed_model(conn: sqlite3.Connection) -> str | None:
|
||||||
|
return IndexMetaTable._get(conn, "embed_model")
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def set_embed_model(conn: sqlite3.Connection, name: str) -> None:
|
||||||
|
IndexMetaTable._set(conn, "embed_model", name)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def get_schema_version(conn: sqlite3.Connection) -> int | None:
|
||||||
|
value = IndexMetaTable._get(conn, "schema_version")
|
||||||
|
return int(value) if value is not None else None
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def set_schema_version(conn: sqlite3.Connection, version: int) -> None:
|
||||||
|
IndexMetaTable._set(conn, "schema_version", str(version))
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def get_total_inserts(conn: sqlite3.Connection) -> int:
|
||||||
|
value = IndexMetaTable._get(conn, "total_inserts")
|
||||||
|
return int(value) if value is not None else 0
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def increment_total_inserts(conn: sqlite3.Connection, count: int) -> None:
|
||||||
|
"""Add ``count`` to the stored counter in one SQL statement (INSERT
|
||||||
|
.. ON CONFLICT DO UPDATE with arithmetic), instead of a separate
|
||||||
|
read-then-write -- called once per add()/upsert_document(), so
|
||||||
|
halving the statement count here is a real, if small, per-call
|
||||||
|
saving. This only avoids a read-then-write race within this single
|
||||||
|
statement; it does not make the counter safe against concurrent
|
||||||
|
writers in general (callers still rely on the write FileLock for
|
||||||
|
that). index_meta.value has TEXT affinity, so the incremented
|
||||||
|
result is stored as its text representation -- get_total_inserts()
|
||||||
|
already expects that (int(value)), so this is not a behavior
|
||||||
|
change, only fewer statements.
|
||||||
|
"""
|
||||||
|
conn.execute(
|
||||||
|
"INSERT INTO index_meta (key, value) VALUES ('total_inserts', ?) "
|
||||||
|
"ON CONFLICT(key) DO UPDATE SET value = "
|
||||||
|
"CAST(index_meta.value AS INTEGER) + CAST(excluded.value AS INTEGER)",
|
||||||
|
(str(count),),
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def reset_total_inserts(conn: sqlite3.Connection, count: int) -> None:
|
||||||
|
"""Set total_inserts to an absolute value -- distinct from
|
||||||
|
increment_total_inserts(): used by compact()'s rebuild and by
|
||||||
|
m0001_v1_to_v2 after copying live rows into a fresh file, where
|
||||||
|
total_inserts must become exactly the live row count, not add to
|
||||||
|
whatever the source file's counter held."""
|
||||||
|
IndexMetaTable._set(conn, "total_inserts", str(count))
|
||||||
@@ -9,6 +9,7 @@ from django.db import connection
|
|||||||
from django.test import override_settings
|
from django.test import override_settings
|
||||||
from django.test.utils import CaptureQueriesContext
|
from django.test.utils import CaptureQueriesContext
|
||||||
from django.utils import timezone
|
from django.utils import timezone
|
||||||
|
from filelock import Timeout
|
||||||
from llama_index.core.schema import MetadataMode
|
from llama_index.core.schema import MetadataMode
|
||||||
|
|
||||||
from documents.models import Correspondent
|
from documents.models import Correspondent
|
||||||
@@ -187,6 +188,45 @@ def test_truncate_embedding_query_returns_single_chunk() -> None:
|
|||||||
assert "word199" not in result
|
assert "word199" not in result
|
||||||
|
|
||||||
|
|
||||||
|
class TestStreamedDocuments:
|
||||||
|
"""_StreamedDocuments streams via .iterator() instead of materializing
|
||||||
|
the whole queryset (plus its content and prefetch caches) in memory at
|
||||||
|
once, while still supporting len() so a progress bar wrapped around it
|
||||||
|
shows a real total.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def test_len_and_iter_delegate_to_streaming_queryset_methods(
|
||||||
|
self,
|
||||||
|
mocker: pytest_mock.MockerFixture,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- A mock queryset
|
||||||
|
WHEN:
|
||||||
|
- A _StreamedDocuments wrapping it is measured and iterated
|
||||||
|
THEN:
|
||||||
|
- len() uses count() (not a materializing len()), and iteration
|
||||||
|
uses .iterator(chunk_size=...) (not plain iteration, which
|
||||||
|
would materialize prefetches for the whole queryset at once)
|
||||||
|
"""
|
||||||
|
mock_queryset = mocker.MagicMock()
|
||||||
|
mock_queryset.count.return_value = 42
|
||||||
|
mock_queryset.iterator.return_value = iter(["doc-1", "doc-2"])
|
||||||
|
streamed = indexing._StreamedDocuments(mock_queryset)
|
||||||
|
|
||||||
|
assert len(streamed) == 42
|
||||||
|
assert list(streamed) == ["doc-1", "doc-2"]
|
||||||
|
# count.call_count isn't asserted exactly: list()'s own size-hint
|
||||||
|
# optimization calls len(streamed) again internally, on top of the
|
||||||
|
# explicit len() call above -- both legitimately delegate to
|
||||||
|
# count(), so only the delegation itself (not the call count) is
|
||||||
|
# the thing being verified here.
|
||||||
|
mock_queryset.count.assert_called_with()
|
||||||
|
mock_queryset.iterator.assert_called_once_with(
|
||||||
|
chunk_size=indexing._INDEX_STREAM_CHUNK_SIZE,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.django_db
|
@pytest.mark.django_db
|
||||||
def test_update_llm_index(
|
def test_update_llm_index(
|
||||||
temp_llm_index_dir: Path,
|
temp_llm_index_dir: Path,
|
||||||
@@ -253,6 +293,45 @@ def test_update_llm_index_rebuilds_on_model_name_change(
|
|||||||
assert store.stored_model_name() == "model-b"
|
assert store.stored_model_name() == "model-b"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.django_db
|
||||||
|
def test_update_llm_index_merges_exists_and_config_mismatch_reads(
|
||||||
|
temp_llm_index_dir: Path,
|
||||||
|
real_document: Document,
|
||||||
|
mock_embed_model: FakeEmbedding,
|
||||||
|
) -> None:
|
||||||
|
# Build an initial index so the second call's table_exists()/
|
||||||
|
# config_mismatch() checks have something real to check against.
|
||||||
|
with patch("documents.models.Document.objects.all") as mock_all:
|
||||||
|
mock_queryset = MagicMock()
|
||||||
|
mock_queryset.exists.return_value = True
|
||||||
|
mock_queryset.__iter__.return_value = iter([real_document])
|
||||||
|
mock_queryset.select_related.return_value = mock_queryset
|
||||||
|
mock_queryset.prefetch_related.return_value = mock_queryset
|
||||||
|
mock_all.return_value = mock_queryset
|
||||||
|
indexing.update_llm_index(rebuild=True)
|
||||||
|
|
||||||
|
with patch("documents.models.Document.objects.all") as mock_all:
|
||||||
|
mock_queryset = MagicMock()
|
||||||
|
mock_queryset.exists.return_value = True
|
||||||
|
mock_queryset.__iter__.return_value = iter([real_document])
|
||||||
|
mock_queryset.select_related.return_value = mock_queryset
|
||||||
|
mock_queryset.prefetch_related.return_value = mock_queryset
|
||||||
|
mock_all.return_value = mock_queryset
|
||||||
|
with patch(
|
||||||
|
"paperless_ai.indexing.read_store",
|
||||||
|
wraps=indexing.read_store,
|
||||||
|
) as read_store_spy:
|
||||||
|
indexing.update_llm_index(rebuild=False)
|
||||||
|
|
||||||
|
# Documents exist, so the fast-exit check's `no_documents and ...`
|
||||||
|
# short-circuits before ever calling llm_index_exists() -- the only
|
||||||
|
# read_store() call left in this path is the merged table_exists()/
|
||||||
|
# config_mismatch() check. Before this task's fix, that merged check
|
||||||
|
# was two separate read_store() calls (one inside llm_index_exists(),
|
||||||
|
# one for config_mismatch() right after) -- so this asserts 1, not 2.
|
||||||
|
assert read_store_spy.call_count == 1
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.django_db
|
@pytest.mark.django_db
|
||||||
def test_update_llm_index_partial_update(
|
def test_update_llm_index_partial_update(
|
||||||
temp_llm_index_dir: Path,
|
temp_llm_index_dir: Path,
|
||||||
@@ -791,6 +870,44 @@ class TestLlmIndexLocking:
|
|||||||
|
|
||||||
mock_store.upsert_document.assert_not_called()
|
mock_store.upsert_document.assert_not_called()
|
||||||
|
|
||||||
|
def test_add_or_update_document_skips_write_when_migration_check_deferred(
|
||||||
|
self,
|
||||||
|
temp_llm_index_dir: Path,
|
||||||
|
mock_embed_model: FakeEmbedding,
|
||||||
|
mocker: pytest_mock.MockerFixture,
|
||||||
|
) -> None:
|
||||||
|
"""A migration check that times out waiting for readers to drain
|
||||||
|
must be treated the same as a pending migration -- proceeding to
|
||||||
|
write would target a store still on its old schema. Regression
|
||||||
|
test for the tri-state fix: a bare bool collapsed this outcome
|
||||||
|
into the same falsy value as "already current".
|
||||||
|
"""
|
||||||
|
mock_store = MagicMock()
|
||||||
|
mock_store.has_pending_migration.return_value = True
|
||||||
|
mocker.patch(
|
||||||
|
"paperless_ai.indexing.write_store",
|
||||||
|
return_value=mocker.MagicMock(
|
||||||
|
__enter__=mocker.MagicMock(return_value=mock_store),
|
||||||
|
__exit__=mocker.MagicMock(return_value=False),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
mocker.patch(
|
||||||
|
"paperless_ai.indexing._exclude_readers",
|
||||||
|
side_effect=Timeout("test"),
|
||||||
|
)
|
||||||
|
mock_node = MagicMock()
|
||||||
|
mock_node.get_content.return_value = "fake node text"
|
||||||
|
mocker.patch(
|
||||||
|
"paperless_ai.indexing.build_document_node",
|
||||||
|
return_value=[mock_node],
|
||||||
|
)
|
||||||
|
|
||||||
|
doc = MagicMock(spec=Document)
|
||||||
|
doc.id = 1
|
||||||
|
indexing.llm_index_add_or_update_document(doc)
|
||||||
|
|
||||||
|
mock_store.upsert_document.assert_not_called()
|
||||||
|
|
||||||
def test_remove_document_uses_write_store(
|
def test_remove_document_uses_write_store(
|
||||||
self,
|
self,
|
||||||
temp_llm_index_dir: Path,
|
temp_llm_index_dir: Path,
|
||||||
@@ -837,6 +954,34 @@ class TestLlmIndexLocking:
|
|||||||
|
|
||||||
mock_store.delete.assert_not_called()
|
mock_store.delete.assert_not_called()
|
||||||
|
|
||||||
|
def test_remove_document_skips_write_when_migration_check_deferred(
|
||||||
|
self,
|
||||||
|
temp_llm_index_dir: Path,
|
||||||
|
mocker: pytest_mock.MockerFixture,
|
||||||
|
) -> None:
|
||||||
|
"""A migration check deferred by a reader-lock timeout must block
|
||||||
|
the delete too, for the same reason as the incremental-update path.
|
||||||
|
"""
|
||||||
|
mock_store = MagicMock()
|
||||||
|
mock_store.has_pending_migration.return_value = True
|
||||||
|
mocker.patch(
|
||||||
|
"paperless_ai.indexing.write_store",
|
||||||
|
return_value=mocker.MagicMock(
|
||||||
|
__enter__=mocker.MagicMock(return_value=mock_store),
|
||||||
|
__exit__=mocker.MagicMock(return_value=False),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
mocker.patch(
|
||||||
|
"paperless_ai.indexing._exclude_readers",
|
||||||
|
side_effect=Timeout("test"),
|
||||||
|
)
|
||||||
|
|
||||||
|
doc = MagicMock(spec=Document)
|
||||||
|
doc.id = 1
|
||||||
|
indexing.llm_index_remove_document(doc)
|
||||||
|
|
||||||
|
mock_store.delete.assert_not_called()
|
||||||
|
|
||||||
def test_update_llm_index_rebuild_uses_write_store(
|
def test_update_llm_index_rebuild_uses_write_store(
|
||||||
self,
|
self,
|
||||||
temp_llm_index_dir: Path,
|
temp_llm_index_dir: Path,
|
||||||
@@ -860,6 +1005,35 @@ class TestLlmIndexLocking:
|
|||||||
|
|
||||||
mock_store.drop_table.assert_called_once()
|
mock_store.drop_table.assert_called_once()
|
||||||
|
|
||||||
|
def test_update_llm_index_skips_when_migration_check_deferred(
|
||||||
|
self,
|
||||||
|
temp_llm_index_dir: Path,
|
||||||
|
mocker: pytest_mock.MockerFixture,
|
||||||
|
) -> None:
|
||||||
|
"""A migration check deferred by a reader-lock timeout must short-
|
||||||
|
circuit before the second write_store() block (document scanning,
|
||||||
|
add/upsert, compaction) ever runs -- that block would otherwise
|
||||||
|
write against a store still on its old schema.
|
||||||
|
"""
|
||||||
|
mock_store = MagicMock()
|
||||||
|
mock_store.has_pending_migration.return_value = True
|
||||||
|
write_store_mock = mocker.patch(
|
||||||
|
"paperless_ai.indexing.write_store",
|
||||||
|
return_value=mocker.MagicMock(
|
||||||
|
__enter__=mocker.MagicMock(return_value=mock_store),
|
||||||
|
__exit__=mocker.MagicMock(return_value=False),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
mocker.patch(
|
||||||
|
"paperless_ai.indexing._exclude_readers",
|
||||||
|
side_effect=Timeout("test"),
|
||||||
|
)
|
||||||
|
|
||||||
|
result = indexing.update_llm_index(rebuild=False)
|
||||||
|
|
||||||
|
assert "deferred" in result
|
||||||
|
write_store_mock.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.django_db
|
@pytest.mark.django_db
|
||||||
@pytest.mark.django_db
|
@pytest.mark.django_db
|
||||||
@@ -978,6 +1152,37 @@ class TestLlmIndexMigrate:
|
|||||||
indexing.llm_index_migrate()
|
indexing.llm_index_migrate()
|
||||||
assert "requires re-embedding" in caplog.text
|
assert "requires re-embedding" in caplog.text
|
||||||
|
|
||||||
|
def test_logs_info_when_migration_check_deferred(
|
||||||
|
self,
|
||||||
|
mocker: pytest_mock.MockerFixture,
|
||||||
|
caplog: pytest.LogCaptureFixture,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- AI/LLM index support is enabled
|
||||||
|
- A pending migration cannot run because readers are active
|
||||||
|
WHEN:
|
||||||
|
- llm_index_migrate() is called
|
||||||
|
THEN:
|
||||||
|
- An info line notes the deferral, not the re-embed warning
|
||||||
|
"""
|
||||||
|
mocker.patch(
|
||||||
|
"paperless_ai.indexing.AIConfig",
|
||||||
|
return_value=mocker.Mock(llm_index_enabled=True),
|
||||||
|
)
|
||||||
|
store_mock = mocker.MagicMock()
|
||||||
|
store_mock.has_pending_migration.return_value = True
|
||||||
|
write_store_cm = mocker.patch("paperless_ai.indexing.write_store")
|
||||||
|
write_store_cm.return_value.__enter__.return_value = store_mock
|
||||||
|
mocker.patch(
|
||||||
|
"paperless_ai.indexing._exclude_readers",
|
||||||
|
side_effect=Timeout("test"),
|
||||||
|
)
|
||||||
|
with caplog.at_level(logging.INFO, logger="paperless_ai.indexing"):
|
||||||
|
indexing.llm_index_migrate()
|
||||||
|
assert "deferred" in caplog.text
|
||||||
|
assert "requires re-embedding" not in caplog.text
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.django_db
|
@pytest.mark.django_db
|
||||||
class TestQuerySimilarDocuments:
|
class TestQuerySimilarDocuments:
|
||||||
|
|||||||
@@ -0,0 +1,340 @@
|
|||||||
|
import sqlite3
|
||||||
|
from collections.abc import Generator
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pytest_mock import MockerFixture
|
||||||
|
|
||||||
|
from paperless_ai.tables import ChunkRow
|
||||||
|
from paperless_ai.tables import DocumentChunksTable
|
||||||
|
from paperless_ai.tables import DocumentMetaRow
|
||||||
|
from paperless_ai.tables import DocumentMetaTable
|
||||||
|
from paperless_ai.tables import IndexMetaTable
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def conn() -> Generator[sqlite3.Connection, None, None]:
|
||||||
|
connection = sqlite3.connect(":memory:")
|
||||||
|
connection.row_factory = sqlite3.Row
|
||||||
|
try:
|
||||||
|
yield connection
|
||||||
|
finally:
|
||||||
|
connection.close()
|
||||||
|
|
||||||
|
|
||||||
|
class TestDocumentChunksTable:
|
||||||
|
def test_create_is_idempotent(self, conn: sqlite3.Connection) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- A bare sqlite3 connection
|
||||||
|
WHEN:
|
||||||
|
- create() is called, a row is inserted, then create() is called again
|
||||||
|
THEN:
|
||||||
|
- No error is raised and the row survives uncorrupted
|
||||||
|
"""
|
||||||
|
DocumentChunksTable.create(conn)
|
||||||
|
DocumentChunksTable.insert_many(conn, [ChunkRow("c1", 1)])
|
||||||
|
DocumentChunksTable.create(conn)
|
||||||
|
assert DocumentChunksTable.chunk_ids_for_document(conn, 1) == ["c1"]
|
||||||
|
|
||||||
|
def test_insert_many_then_lookup_by_document_id(
|
||||||
|
self,
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- An empty document_chunks table
|
||||||
|
WHEN:
|
||||||
|
- Two chunks for document 1 and one for document 2 are inserted
|
||||||
|
THEN:
|
||||||
|
- chunk_ids_for_document returns exactly the matching chunk ids
|
||||||
|
"""
|
||||||
|
DocumentChunksTable.create(conn)
|
||||||
|
DocumentChunksTable.insert_many(
|
||||||
|
conn,
|
||||||
|
[ChunkRow("c1", 1), ChunkRow("c2", 1), ChunkRow("c3", 2)],
|
||||||
|
)
|
||||||
|
assert sorted(DocumentChunksTable.chunk_ids_for_document(conn, 1)) == [
|
||||||
|
"c1",
|
||||||
|
"c2",
|
||||||
|
]
|
||||||
|
assert DocumentChunksTable.chunk_ids_for_document(conn, 2) == ["c3"]
|
||||||
|
assert DocumentChunksTable.chunk_ids_for_document(conn, 999) == []
|
||||||
|
|
||||||
|
def test_delete_for_document_removes_only_that_document(
|
||||||
|
self,
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- Chunks for two different documents
|
||||||
|
WHEN:
|
||||||
|
- delete_for_document() is called for one of them
|
||||||
|
THEN:
|
||||||
|
- Only that document's chunks are removed
|
||||||
|
"""
|
||||||
|
DocumentChunksTable.create(conn)
|
||||||
|
DocumentChunksTable.insert_many(
|
||||||
|
conn,
|
||||||
|
[ChunkRow("c1", 1), ChunkRow("c2", 2)],
|
||||||
|
)
|
||||||
|
DocumentChunksTable.delete_for_document(conn, 1)
|
||||||
|
assert DocumentChunksTable.chunk_ids_for_document(conn, 1) == []
|
||||||
|
assert DocumentChunksTable.chunk_ids_for_document(conn, 2) == ["c2"]
|
||||||
|
|
||||||
|
def test_delete_all_clears_every_row(self, conn: sqlite3.Connection) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- Chunks for multiple documents
|
||||||
|
WHEN:
|
||||||
|
- delete_all() is called
|
||||||
|
THEN:
|
||||||
|
- count() returns 0
|
||||||
|
"""
|
||||||
|
DocumentChunksTable.create(conn)
|
||||||
|
DocumentChunksTable.insert_many(
|
||||||
|
conn,
|
||||||
|
[ChunkRow("c1", 1), ChunkRow("c2", 2)],
|
||||||
|
)
|
||||||
|
DocumentChunksTable.delete_all(conn)
|
||||||
|
assert DocumentChunksTable.count(conn) == 0
|
||||||
|
|
||||||
|
def test_count_reflects_live_rows(self, conn: sqlite3.Connection) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- An empty document_chunks table
|
||||||
|
WHEN:
|
||||||
|
- Rows are inserted then one document's rows are deleted
|
||||||
|
THEN:
|
||||||
|
- count() reflects the remaining row count
|
||||||
|
"""
|
||||||
|
DocumentChunksTable.create(conn)
|
||||||
|
DocumentChunksTable.insert_many(
|
||||||
|
conn,
|
||||||
|
[ChunkRow("c1", 1), ChunkRow("c2", 1), ChunkRow("c3", 2)],
|
||||||
|
)
|
||||||
|
assert DocumentChunksTable.count(conn) == 3
|
||||||
|
DocumentChunksTable.delete_for_document(conn, 1)
|
||||||
|
assert DocumentChunksTable.count(conn) == 1
|
||||||
|
|
||||||
|
|
||||||
|
class TestDocumentMetaTable:
|
||||||
|
def test_upsert_many_then_all_modified_times(
|
||||||
|
self,
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- An empty document_meta table
|
||||||
|
WHEN:
|
||||||
|
- Two documents' modified timestamps are upserted
|
||||||
|
THEN:
|
||||||
|
- all_modified_times() returns both, keyed by str(document_id)
|
||||||
|
"""
|
||||||
|
DocumentMetaTable.create(conn)
|
||||||
|
DocumentMetaTable.upsert_many(
|
||||||
|
conn,
|
||||||
|
[
|
||||||
|
DocumentMetaRow(1, "2026-01-01T00:00:00"),
|
||||||
|
DocumentMetaRow(2, "2026-02-02T00:00:00"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
assert DocumentMetaTable.all_modified_times(conn) == {
|
||||||
|
"1": "2026-01-01T00:00:00",
|
||||||
|
"2": "2026-02-02T00:00:00",
|
||||||
|
}
|
||||||
|
|
||||||
|
def test_upsert_many_overwrites_existing_value(
|
||||||
|
self,
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- A document_meta row for document 1
|
||||||
|
WHEN:
|
||||||
|
- upsert_many() is called again with a new modified value for
|
||||||
|
the same document_id
|
||||||
|
THEN:
|
||||||
|
- The stored value is replaced, not duplicated
|
||||||
|
"""
|
||||||
|
DocumentMetaTable.create(conn)
|
||||||
|
DocumentMetaTable.upsert_many(conn, [DocumentMetaRow(1, "old")])
|
||||||
|
DocumentMetaTable.upsert_many(conn, [DocumentMetaRow(1, "new")])
|
||||||
|
assert DocumentMetaTable.all_modified_times(conn) == {"1": "new"}
|
||||||
|
|
||||||
|
def test_delete_for_document_removes_only_that_row(
|
||||||
|
self,
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- document_meta rows for two documents
|
||||||
|
WHEN:
|
||||||
|
- delete_for_document() is called for one of them
|
||||||
|
THEN:
|
||||||
|
- Only that document's row is removed
|
||||||
|
"""
|
||||||
|
DocumentMetaTable.create(conn)
|
||||||
|
DocumentMetaTable.upsert_many(
|
||||||
|
conn,
|
||||||
|
[DocumentMetaRow(1, "a"), DocumentMetaRow(2, "b")],
|
||||||
|
)
|
||||||
|
DocumentMetaTable.delete_for_document(conn, 1)
|
||||||
|
assert DocumentMetaTable.all_modified_times(conn) == {"2": "b"}
|
||||||
|
|
||||||
|
def test_delete_all_clears_every_row(self, conn: sqlite3.Connection) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- document_meta rows for multiple documents
|
||||||
|
WHEN:
|
||||||
|
- delete_all() is called
|
||||||
|
THEN:
|
||||||
|
- all_modified_times() returns an empty dict
|
||||||
|
"""
|
||||||
|
DocumentMetaTable.create(conn)
|
||||||
|
DocumentMetaTable.upsert_many(
|
||||||
|
conn,
|
||||||
|
[DocumentMetaRow(1, "a"), DocumentMetaRow(2, "b")],
|
||||||
|
)
|
||||||
|
DocumentMetaTable.delete_all(conn)
|
||||||
|
assert DocumentMetaTable.all_modified_times(conn) == {}
|
||||||
|
|
||||||
|
def test_copy_all_streams_every_row_to_destination(
|
||||||
|
self,
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- A source connection with document_meta rows for 5 documents
|
||||||
|
- A separate, empty destination connection
|
||||||
|
WHEN:
|
||||||
|
- copy_all() is called with a batch size smaller than the row
|
||||||
|
count, forcing multiple fetchmany() cycles
|
||||||
|
THEN:
|
||||||
|
- Every row is present on the destination connection
|
||||||
|
"""
|
||||||
|
DocumentMetaTable.create(conn)
|
||||||
|
DocumentMetaTable.upsert_many(
|
||||||
|
conn,
|
||||||
|
[DocumentMetaRow(i, f"modified-{i}") for i in range(5)],
|
||||||
|
)
|
||||||
|
dst_conn = sqlite3.connect(":memory:")
|
||||||
|
dst_conn.row_factory = sqlite3.Row
|
||||||
|
try:
|
||||||
|
DocumentMetaTable.create(dst_conn)
|
||||||
|
DocumentMetaTable.copy_all(conn, dst_conn, batch_size=2)
|
||||||
|
assert DocumentMetaTable.all_modified_times(dst_conn) == {
|
||||||
|
str(i): f"modified-{i}" for i in range(5)
|
||||||
|
}
|
||||||
|
finally:
|
||||||
|
dst_conn.close()
|
||||||
|
|
||||||
|
|
||||||
|
class TestIndexMetaTable:
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("setter_name", "getter_name", "value"),
|
||||||
|
[
|
||||||
|
("set_dim", "get_dim", 384),
|
||||||
|
("set_embed_model", "get_embed_model", "model-a"),
|
||||||
|
("set_schema_version", "get_schema_version", 2),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_typed_accessor_roundtrip(
|
||||||
|
self,
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
setter_name: str,
|
||||||
|
getter_name: str,
|
||||||
|
value: int | str,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- An empty index_meta table
|
||||||
|
WHEN:
|
||||||
|
- A typed accessor's setter is called then the getter is read back
|
||||||
|
THEN:
|
||||||
|
- The same value is returned, correctly typed (int or str)
|
||||||
|
"""
|
||||||
|
IndexMetaTable.create(conn)
|
||||||
|
getter = getattr(IndexMetaTable, getter_name)
|
||||||
|
setter = getattr(IndexMetaTable, setter_name)
|
||||||
|
assert getter(conn) is None
|
||||||
|
setter(conn, value)
|
||||||
|
assert getter(conn) == value
|
||||||
|
|
||||||
|
def test_total_inserts_starts_at_zero(self, conn: sqlite3.Connection) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- An empty index_meta table
|
||||||
|
WHEN:
|
||||||
|
- get_total_inserts() is read before anything is set
|
||||||
|
THEN:
|
||||||
|
- 0 is returned
|
||||||
|
"""
|
||||||
|
IndexMetaTable.create(conn)
|
||||||
|
assert IndexMetaTable.get_total_inserts(conn) == 0
|
||||||
|
|
||||||
|
def test_increment_total_inserts_accumulates(
|
||||||
|
self,
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- An empty index_meta table
|
||||||
|
WHEN:
|
||||||
|
- increment_total_inserts() is called twice
|
||||||
|
THEN:
|
||||||
|
- get_total_inserts() returns the running sum
|
||||||
|
"""
|
||||||
|
IndexMetaTable.create(conn)
|
||||||
|
IndexMetaTable.increment_total_inserts(conn, 5)
|
||||||
|
IndexMetaTable.increment_total_inserts(conn, 3)
|
||||||
|
assert IndexMetaTable.get_total_inserts(conn) == 8
|
||||||
|
|
||||||
|
def test_increment_total_inserts_is_a_single_statement(
|
||||||
|
self,
|
||||||
|
mocker: MockerFixture,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- An empty index_meta table
|
||||||
|
WHEN:
|
||||||
|
- increment_total_inserts() is called
|
||||||
|
THEN:
|
||||||
|
- Exactly one conn.execute() call is made (a single INSERT ...
|
||||||
|
ON CONFLICT DO UPDATE, not a separate read then write)
|
||||||
|
"""
|
||||||
|
|
||||||
|
# sqlite3.Connection is an immutable C extension type with no
|
||||||
|
# instance __dict__, so mocker.spy(conn, "execute") can't shadow
|
||||||
|
# "execute" on a plain connection ("attribute 'execute' is
|
||||||
|
# read-only"). A trivial Python subclass gets a normal instance
|
||||||
|
# __dict__, making the instance spyable while still being a real,
|
||||||
|
# usable sqlite3.Connection.
|
||||||
|
class _SpyableConnection(sqlite3.Connection):
|
||||||
|
pass
|
||||||
|
|
||||||
|
conn = sqlite3.connect(":memory:", factory=_SpyableConnection)
|
||||||
|
try:
|
||||||
|
conn.row_factory = sqlite3.Row
|
||||||
|
IndexMetaTable.create(conn)
|
||||||
|
execute_spy = mocker.spy(conn, "execute")
|
||||||
|
IndexMetaTable.increment_total_inserts(conn, 5)
|
||||||
|
assert execute_spy.call_count == 1
|
||||||
|
finally:
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
def test_reset_total_inserts_sets_absolute_value(
|
||||||
|
self,
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- A total_inserts counter already at a high value
|
||||||
|
WHEN:
|
||||||
|
- reset_total_inserts() is called with a lower value
|
||||||
|
THEN:
|
||||||
|
- get_total_inserts() returns exactly that value, not a sum
|
||||||
|
"""
|
||||||
|
IndexMetaTable.create(conn)
|
||||||
|
IndexMetaTable.increment_total_inserts(conn, 100)
|
||||||
|
IndexMetaTable.reset_total_inserts(conn, 7)
|
||||||
|
assert IndexMetaTable.get_total_inserts(conn) == 7
|
||||||
@@ -1,28 +1,35 @@
|
|||||||
|
import inspect
|
||||||
import sqlite3
|
import sqlite3
|
||||||
from collections.abc import Generator
|
from collections.abc import Generator
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
import sqlite_vec
|
||||||
from llama_index.core.schema import TextNode
|
from llama_index.core.schema import TextNode
|
||||||
from llama_index.core.vector_stores.types import FilterOperator
|
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 MetadataFilter
|
||||||
from llama_index.core.vector_stores.types import MetadataFilters
|
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 VectorStoreQuery
|
||||||
|
from pytest_mock import MockerFixture
|
||||||
|
|
||||||
from paperless_ai.migrations import MIGRATIONS
|
from paperless_ai.migrations import MIGRATIONS
|
||||||
from paperless_ai.migrations import Migration
|
from paperless_ai.migrations import Migration
|
||||||
|
from paperless_ai.migrations import m0001_v1_to_v2
|
||||||
|
from paperless_ai.tables import DocumentChunksTable
|
||||||
|
from paperless_ai.tables import DocumentMetaTable
|
||||||
from paperless_ai.vector_store import DB_FILENAME
|
from paperless_ai.vector_store import DB_FILENAME
|
||||||
from paperless_ai.vector_store import DEFAULT_TABLE_NAME
|
from paperless_ai.vector_store import DEFAULT_TABLE_NAME
|
||||||
from paperless_ai.vector_store import SCHEMA_VERSION
|
from paperless_ai.vector_store import SCHEMA_VERSION
|
||||||
from paperless_ai.vector_store import PaperlessSqliteVecVectorStore
|
from paperless_ai.vector_store import PaperlessSqliteVecVectorStore
|
||||||
from paperless_ai.vector_store import _build_where
|
from paperless_ai.vector_store import _build_where
|
||||||
|
from paperless_ai.vector_store import _pack
|
||||||
|
|
||||||
DIM = 16
|
DIM = 16
|
||||||
|
|
||||||
|
|
||||||
def make_node(
|
def make_node(
|
||||||
node_id: str,
|
node_id: str,
|
||||||
document_id: str,
|
document_id: int,
|
||||||
*,
|
*,
|
||||||
modified: str = "2026-06-10T00:00:00",
|
modified: str = "2026-06-10T00:00:00",
|
||||||
seed: float = 0.0,
|
seed: float = 0.0,
|
||||||
@@ -59,13 +66,13 @@ def _query(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _eq_filter(key: str, value: str):
|
def _eq_filter(key: str, value: int):
|
||||||
return MetadataFilters(
|
return MetadataFilters(
|
||||||
filters=[MetadataFilter(key=key, operator=FilterOperator.EQ, value=value)],
|
filters=[MetadataFilter(key=key, operator=FilterOperator.EQ, value=value)],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _in_filter(document_ids: list[str]):
|
def _in_filter(document_ids: list[int]):
|
||||||
return MetadataFilters(
|
return MetadataFilters(
|
||||||
filters=[
|
filters=[
|
||||||
MetadataFilter(
|
MetadataFilter(
|
||||||
@@ -77,7 +84,7 @@ def _in_filter(document_ids: list[str]):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _ne_filter(document_id: str):
|
def _ne_filter(document_id: int):
|
||||||
return MetadataFilters(
|
return MetadataFilters(
|
||||||
filters=[
|
filters=[
|
||||||
MetadataFilter(
|
MetadataFilter(
|
||||||
@@ -91,11 +98,11 @@ def _ne_filter(document_id: str):
|
|||||||
|
|
||||||
class TestCrud:
|
class TestCrud:
|
||||||
def test_add_then_query_returns_node(self, store) -> None:
|
def test_add_then_query_returns_node(self, store) -> None:
|
||||||
node = make_node("n1", "1")
|
node = make_node("n1", 1)
|
||||||
assert store.add([node]) == ["n1"]
|
assert store.add([node]) == ["n1"]
|
||||||
result = _query(store, node.embedding, top_k=1)
|
result = _query(store, node.embedding, top_k=1)
|
||||||
assert result.ids == ["n1"]
|
assert result.ids == ["n1"]
|
||||||
assert result.nodes[0].metadata["document_id"] == "1"
|
assert result.nodes[0].metadata["document_id"] == 1
|
||||||
# cosine distance of the identical vector is 0 -> similarity 1
|
# cosine distance of the identical vector is 0 -> similarity 1
|
||||||
assert result.similarities[0] == pytest.approx(1.0)
|
assert result.similarities[0] == pytest.approx(1.0)
|
||||||
|
|
||||||
@@ -108,58 +115,58 @@ class TestCrud:
|
|||||||
assert not store.table_exists()
|
assert not store.table_exists()
|
||||||
|
|
||||||
def test_delete_removes_all_chunks_of_document(self, store) -> None:
|
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.add([make_node("a1", 1), make_node("a2", 1), make_node("b1", 2)])
|
||||||
store.delete("1")
|
store.delete(1)
|
||||||
result = _query(store, [0.0] * DIM, top_k=10)
|
result = _query(store, [0.0] * DIM, top_k=10)
|
||||||
assert result.ids == ["b1"]
|
assert result.ids == ["b1"]
|
||||||
|
|
||||||
def test_query_with_in_filter_scopes_results(self, store) -> None:
|
def test_query_with_in_filter_scopes_results(self, store) -> None:
|
||||||
store.add(
|
store.add(
|
||||||
[
|
[
|
||||||
make_node("a1", "1", seed=0.0),
|
make_node("a1", 1, seed=0.0),
|
||||||
make_node("b1", "2", seed=1.0),
|
make_node("b1", 2, seed=1.0),
|
||||||
make_node("c1", "3", seed=2.0),
|
make_node("c1", 3, seed=2.0),
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
result = _query(store, [0.0] * DIM, top_k=10, filters=_in_filter(["2", "3"]))
|
result = _query(store, [0.0] * DIM, top_k=10, filters=_in_filter([2, 3]))
|
||||||
assert sorted(result.ids) == ["b1", "c1"]
|
assert sorted(result.ids) == ["b1", "c1"]
|
||||||
|
|
||||||
def test_query_respects_top_k_with_filter(self, store) -> None:
|
def test_query_respects_top_k_with_filter(self, store) -> None:
|
||||||
# k semantics: global top-k even with IN filters (document_id is a
|
# k semantics: global top-k even with IN filters (document_id is a
|
||||||
# metadata column, not a partition key -- see design doc).
|
# metadata column, not a partition key -- see design doc).
|
||||||
store.add(
|
store.add(
|
||||||
[make_node(f"n{i}", str(i % 4), seed=float(i)) for i in range(12)],
|
[make_node(f"n{i}", i % 4, seed=float(i)) for i in range(12)],
|
||||||
)
|
)
|
||||||
result = _query(
|
result = _query(
|
||||||
store,
|
store,
|
||||||
[0.0] * DIM,
|
[0.0] * DIM,
|
||||||
top_k=3,
|
top_k=3,
|
||||||
filters=_in_filter(["0", "1", "2", "3"]),
|
filters=_in_filter([0, 1, 2, 3]),
|
||||||
)
|
)
|
||||||
assert len(result.ids) == 3
|
assert len(result.ids) == 3
|
||||||
assert result.similarities == sorted(result.similarities, reverse=True)
|
assert result.similarities == sorted(result.similarities, reverse=True)
|
||||||
|
|
||||||
def test_get_nodes_filter_and_empty_paths(self, store) -> None:
|
def test_get_nodes_filter_and_empty_paths(self, store) -> None:
|
||||||
assert store.get_nodes(filters=_in_filter(["1"])) == [] # no table yet
|
assert store.get_nodes(filters=_in_filter([1])) == [] # no table yet
|
||||||
store.add([make_node("a1", "1"), make_node("b1", "2")])
|
store.add([make_node("a1", 1), make_node("b1", 2)])
|
||||||
nodes = store.get_nodes(filters=_in_filter(["1"]))
|
nodes = store.get_nodes(filters=_in_filter([1]))
|
||||||
assert [n.node_id for n in nodes] == ["a1"]
|
assert [n.node_id for n in nodes] == ["a1"]
|
||||||
assert nodes[0].embedding is not None
|
assert nodes[0].embedding is not None
|
||||||
assert store.get_nodes(filters=_in_filter(["999"])) == []
|
assert store.get_nodes(filters=_in_filter([999])) == []
|
||||||
|
|
||||||
def test_query_with_eq_filter_scopes_results(self, store) -> None:
|
def test_query_with_eq_filter_scopes_results(self, store) -> None:
|
||||||
store.add(
|
store.add(
|
||||||
[
|
[
|
||||||
make_node("a1", "1", seed=0.0),
|
make_node("a1", 1, seed=0.0),
|
||||||
make_node("b1", "2", seed=1.0),
|
make_node("b1", 2, seed=1.0),
|
||||||
make_node("c1", "3", seed=2.0),
|
make_node("c1", 3, seed=2.0),
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
result = _query(
|
result = _query(
|
||||||
store,
|
store,
|
||||||
[0.0] * DIM,
|
[0.0] * DIM,
|
||||||
top_k=10,
|
top_k=10,
|
||||||
filters=_eq_filter("document_id", "2"),
|
filters=_eq_filter("document_id", 2),
|
||||||
)
|
)
|
||||||
assert result.ids == ["b1"]
|
assert result.ids == ["b1"]
|
||||||
|
|
||||||
@@ -168,7 +175,7 @@ class TestCrud:
|
|||||||
store.get_nodes(node_ids=["x"])
|
store.get_nodes(node_ids=["x"])
|
||||||
|
|
||||||
def test_fresh_instance_sees_existing_table(self, store, tmp_path: Path) -> None:
|
def test_fresh_instance_sees_existing_table(self, store, tmp_path: Path) -> None:
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reopened:
|
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reopened:
|
||||||
assert reopened.table_exists()
|
assert reopened.table_exists()
|
||||||
assert reopened.vector_dim() == DIM
|
assert reopened.vector_dim() == DIM
|
||||||
@@ -176,23 +183,98 @@ class TestCrud:
|
|||||||
|
|
||||||
def test_table_exists_and_drop(self, store) -> None:
|
def test_table_exists_and_drop(self, store) -> None:
|
||||||
assert not store.table_exists()
|
assert not store.table_exists()
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
assert store.table_exists()
|
assert store.table_exists()
|
||||||
store.drop_table()
|
store.drop_table()
|
||||||
assert not store.table_exists()
|
assert not store.table_exists()
|
||||||
assert store.vector_dim() is None
|
assert store.vector_dim() is None
|
||||||
|
|
||||||
|
def test_document_id_stored_as_integer_in_vec0(
|
||||||
|
self,
|
||||||
|
store: PaperlessSqliteVecVectorStore,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- An empty vector store
|
||||||
|
WHEN:
|
||||||
|
- A node is added with an int document_id
|
||||||
|
THEN:
|
||||||
|
- vec0's own document_id column holds an INTEGER, not TEXT
|
||||||
|
"""
|
||||||
|
store.add([make_node("a1", 1)])
|
||||||
|
row = store.client.execute(
|
||||||
|
"SELECT document_id FROM documents WHERE id = 'a1'",
|
||||||
|
).fetchone()
|
||||||
|
assert isinstance(row["document_id"], int)
|
||||||
|
|
||||||
|
def test_drop_table_clears_modified_times(
|
||||||
|
self,
|
||||||
|
store: PaperlessSqliteVecVectorStore,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- A store with a tracked document's modified time
|
||||||
|
WHEN:
|
||||||
|
- drop_table() is called
|
||||||
|
THEN:
|
||||||
|
- document_meta and document_chunks are both cleared directly
|
||||||
|
(asserted against the tables themselves, not via
|
||||||
|
get_modified_times()/table_exists() -- those short-circuit on
|
||||||
|
the vec0 table being gone, which drop_table() does first, so
|
||||||
|
they would pass even if DocumentMetaTable.delete_all()/
|
||||||
|
DocumentChunksTable.delete_all() were never called)
|
||||||
|
"""
|
||||||
|
store.add([make_node("a1", 1)])
|
||||||
|
store.drop_table()
|
||||||
|
assert (
|
||||||
|
store.client.execute(
|
||||||
|
"SELECT count(*) FROM document_meta",
|
||||||
|
).fetchone()[0]
|
||||||
|
== 0
|
||||||
|
)
|
||||||
|
assert (
|
||||||
|
store.client.execute(
|
||||||
|
"SELECT count(*) FROM document_chunks",
|
||||||
|
).fetchone()[0]
|
||||||
|
== 0
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_upsert_document_checks_table_exists_once(
|
||||||
|
self,
|
||||||
|
store: PaperlessSqliteVecVectorStore,
|
||||||
|
mocker: MockerFixture,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- An existing store with one document already indexed
|
||||||
|
WHEN:
|
||||||
|
- upsert_document() replaces that document's chunks
|
||||||
|
THEN:
|
||||||
|
- table_exists() is queried at most once per call, not twice
|
||||||
|
(previously: once via _ensure_table(), once via the separate
|
||||||
|
`if self.table_exists():` delete-chunks guard)
|
||||||
|
"""
|
||||||
|
store.add([make_node("a1", 1)])
|
||||||
|
# store is a pydantic model, whose __setattr__/__delattr__ reject
|
||||||
|
# arbitrary instance attributes ("object has no attribute
|
||||||
|
# 'table_exists'"), so mocker.spy(store, "table_exists") can't
|
||||||
|
# shadow the method on the instance. Spying on the class works
|
||||||
|
# (bound method lookup on the instance still resolves through it).
|
||||||
|
exists_spy = mocker.spy(PaperlessSqliteVecVectorStore, "table_exists")
|
||||||
|
store.upsert_document(1, [make_node("a2", 1)])
|
||||||
|
assert exists_spy.call_count == 1
|
||||||
|
|
||||||
|
|
||||||
class TestBuildWhere:
|
class TestBuildWhere:
|
||||||
def test_ne_filter_translates_to_not_equal_clause(self) -> None:
|
def test_ne_filter_translates_to_not_equal_clause(self) -> None:
|
||||||
where, params = _build_where(_ne_filter("1"))
|
where, params = _build_where(_ne_filter(1))
|
||||||
assert where == "(document_id != ?)"
|
assert where == "(document_id != ?)"
|
||||||
assert params == ["1"]
|
assert params == [1]
|
||||||
|
|
||||||
def test_query_with_ne_filter_excludes_matching_document(self, store) -> None:
|
def test_query_with_ne_filter_excludes_matching_document(self, store) -> None:
|
||||||
store.add([make_node("a1", "1"), make_node("b1", "2")])
|
store.add([make_node("a1", 1), make_node("b1", 2)])
|
||||||
assert sorted(
|
assert sorted(
|
||||||
_query(store, [0.0] * DIM, top_k=5, filters=_ne_filter("1")).ids,
|
_query(store, [0.0] * DIM, top_k=5, filters=_ne_filter(1)).ids,
|
||||||
) == [
|
) == [
|
||||||
"b1",
|
"b1",
|
||||||
]
|
]
|
||||||
@@ -206,7 +288,7 @@ class TestBuildWhere:
|
|||||||
MetadataFilter(
|
MetadataFilter(
|
||||||
key="document_id",
|
key="document_id",
|
||||||
operator=FilterOperator.EQ,
|
operator=FilterOperator.EQ,
|
||||||
value="1",
|
value=1,
|
||||||
),
|
),
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
@@ -215,13 +297,13 @@ class TestBuildWhere:
|
|||||||
assert params == []
|
assert params == []
|
||||||
|
|
||||||
def test_query_with_untranslatable_filter_returns_no_rows(self, store) -> None:
|
def test_query_with_untranslatable_filter_returns_no_rows(self, store) -> None:
|
||||||
store.add([make_node("a1", "1"), make_node("b1", "2")])
|
store.add([make_node("a1", 1), make_node("b1", 2)])
|
||||||
nested = MetadataFilters(
|
nested = MetadataFilters(
|
||||||
filters=[
|
filters=[
|
||||||
MetadataFilter(
|
MetadataFilter(
|
||||||
key="document_id",
|
key="document_id",
|
||||||
operator=FilterOperator.EQ,
|
operator=FilterOperator.EQ,
|
||||||
value="1",
|
value=1,
|
||||||
),
|
),
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
@@ -234,19 +316,19 @@ class TestBuildWhere:
|
|||||||
class TestUpsert:
|
class TestUpsert:
|
||||||
def test_upsert_replaces_and_prunes_stale_chunks(self, store) -> None:
|
def test_upsert_replaces_and_prunes_stale_chunks(self, store) -> None:
|
||||||
store.add(
|
store.add(
|
||||||
[make_node("d1c1", "1"), make_node("d1c2", "1"), make_node("d2c1", "2")],
|
[make_node("d1c1", 1), make_node("d1c2", 1), make_node("d2c1", 2)],
|
||||||
)
|
)
|
||||||
store.upsert_document("1", [make_node("d1new", "1")])
|
store.upsert_document(1, [make_node("d1new", 1)])
|
||||||
result = _query(store, [0.0] * DIM, top_k=10)
|
result = _query(store, [0.0] * DIM, top_k=10)
|
||||||
assert sorted(result.ids) == ["d1new", "d2c1"]
|
assert sorted(result.ids) == ["d1new", "d2c1"]
|
||||||
|
|
||||||
def test_upsert_creates_table_when_missing(self, store) -> None:
|
def test_upsert_creates_table_when_missing(self, store) -> None:
|
||||||
store.upsert_document("1", [make_node("a1", "1")])
|
store.upsert_document(1, [make_node("a1", 1)])
|
||||||
assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"]
|
assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"]
|
||||||
|
|
||||||
def test_upsert_empty_nodes_removes_document(self, store) -> None:
|
def test_upsert_empty_nodes_removes_document(self, store) -> None:
|
||||||
store.add([make_node("a1", "1"), make_node("b1", "2")])
|
store.add([make_node("a1", 1), make_node("b1", 2)])
|
||||||
store.upsert_document("1", [])
|
store.upsert_document(1, [])
|
||||||
assert _query(store, [0.0] * DIM, top_k=10).ids == ["b1"]
|
assert _query(store, [0.0] * DIM, top_k=10).ids == ["b1"]
|
||||||
|
|
||||||
def test_upsert_is_atomic_for_concurrent_readers(
|
def test_upsert_is_atomic_for_concurrent_readers(
|
||||||
@@ -255,16 +337,16 @@ class TestUpsert:
|
|||||||
tmp_path: Path,
|
tmp_path: Path,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""A second connection must never observe document 1 half-replaced."""
|
"""A second connection must never observe document 1 half-replaced."""
|
||||||
store.add([make_node("a1", "1"), make_node("a2", "1")])
|
store.add([make_node("a1", 1), make_node("a2", 1)])
|
||||||
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reader:
|
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reader:
|
||||||
store.upsert_document("1", [make_node("a3", "1")])
|
store.upsert_document(1, [make_node("a3", 1)])
|
||||||
ids = [n.node_id for n in reader.get_nodes(filters=_in_filter(["1"]))]
|
ids = [n.node_id for n in reader.get_nodes(filters=_in_filter([1]))]
|
||||||
assert ids == ["a3"]
|
assert ids == ["a3"]
|
||||||
|
|
||||||
|
|
||||||
class TestMetadataCoercion:
|
class TestMetadataCoercion:
|
||||||
def test_none_metadata_values_become_empty_strings(self, store) -> None:
|
def test_none_metadata_values_become_empty_strings(self, store) -> None:
|
||||||
node = make_node("a1", "1")
|
node = make_node("a1", 1)
|
||||||
node.metadata["modified"] = None
|
node.metadata["modified"] = None
|
||||||
store.add([node]) # must not raise (vec0 rejects NULL metadata)
|
store.add([node]) # must not raise (vec0 rejects NULL metadata)
|
||||||
assert store.get_modified_times() == {"1": ""}
|
assert store.get_modified_times() == {"1": ""}
|
||||||
@@ -283,7 +365,7 @@ class TestModelNameTracking:
|
|||||||
uri=str(tmp_path),
|
uri=str(tmp_path),
|
||||||
embed_model_name="model-a",
|
embed_model_name="model-a",
|
||||||
) as store:
|
) as store:
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
assert store.stored_model_name() == "model-a"
|
assert store.stored_model_name() == "model-a"
|
||||||
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reopened:
|
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reopened:
|
||||||
assert reopened.stored_model_name() == "model-a"
|
assert reopened.stored_model_name() == "model-a"
|
||||||
@@ -294,7 +376,7 @@ class TestModelNameTracking:
|
|||||||
embed_model_name="model-a",
|
embed_model_name="model-a",
|
||||||
) as store:
|
) as store:
|
||||||
assert not store.config_mismatch("anything") # no table yet
|
assert not store.config_mismatch("anything") # no table yet
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
assert not store.config_mismatch("model-a")
|
assert not store.config_mismatch("model-a")
|
||||||
assert store.config_mismatch("model-b")
|
assert store.config_mismatch("model-b")
|
||||||
|
|
||||||
@@ -303,7 +385,7 @@ class TestModelNameTracking:
|
|||||||
tmp_path: Path,
|
tmp_path: Path,
|
||||||
) -> None:
|
) -> None:
|
||||||
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as store: # no model name
|
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as store: # no model name
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
assert not store.config_mismatch("model-a")
|
assert not store.config_mismatch("model-a")
|
||||||
|
|
||||||
|
|
||||||
@@ -314,9 +396,9 @@ class TestGetModifiedTimes:
|
|||||||
def test_returns_one_entry_per_document(self, store) -> None:
|
def test_returns_one_entry_per_document(self, store) -> None:
|
||||||
store.add(
|
store.add(
|
||||||
[
|
[
|
||||||
make_node("a1", "1", modified="2026-01-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("a2", 1, modified="2026-01-01T00:00:00"),
|
||||||
make_node("b1", "2", modified="2026-02-02T00:00:00"),
|
make_node("b1", 2, modified="2026-02-02T00:00:00"),
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
assert store.get_modified_times() == {
|
assert store.get_modified_times() == {
|
||||||
@@ -341,37 +423,35 @@ class TestCompact:
|
|||||||
def _churn(self, store, cycles: int) -> None:
|
def _churn(self, store, cycles: int) -> None:
|
||||||
for i in range(cycles):
|
for i in range(cycles):
|
||||||
store.upsert_document(
|
store.upsert_document(
|
||||||
"1",
|
1,
|
||||||
[make_node(f"gen{i}-{j}", "1", seed=float(j)) for j in range(20)],
|
[make_node(f"gen{i}-{j}", 1, seed=float(j)) for j in range(20)],
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_compact_noop_below_threshold(self, store) -> None:
|
def test_compact_noop_below_threshold(self, store) -> None:
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
store.compact()
|
store.compact()
|
||||||
assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"]
|
assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"]
|
||||||
|
|
||||||
def test_force_compact_preserves_rows_and_metadata(self, store) -> None:
|
def test_force_compact_preserves_rows_and_metadata(self, store) -> None:
|
||||||
store.add([make_node("a1", "1"), make_node("b1", "2", seed=3.0)])
|
store.add([make_node("a1", 1), make_node("b1", 2, seed=3.0)])
|
||||||
self._churn(store, 5)
|
self._churn(store, 5)
|
||||||
before = {
|
before = {
|
||||||
n.node_id: n.metadata
|
n.node_id: n.metadata for n in store.get_nodes(filters=_in_filter([1, 2]))
|
||||||
for n in store.get_nodes(filters=_in_filter(["1", "2"]))
|
|
||||||
}
|
}
|
||||||
store.compact(force=True)
|
store.compact(force=True)
|
||||||
after = {
|
after = {
|
||||||
n.node_id: n.metadata
|
n.node_id: n.metadata for n in store.get_nodes(filters=_in_filter([1, 2]))
|
||||||
for n in store.get_nodes(filters=_in_filter(["1", "2"]))
|
|
||||||
}
|
}
|
||||||
assert after == before
|
assert after == before
|
||||||
assert self._bloat_ratio(store) == pytest.approx(1.0)
|
assert self._bloat_ratio(store) == pytest.approx(1.0)
|
||||||
# store remains fully usable after the rebuild; use a seed far from all
|
# 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
|
# existing nodes (gen4-0..gen4-19 have seeds 0..19) so cosine KNN is
|
||||||
# unambiguous at top_k=1.
|
# unambiguous at top_k=1.
|
||||||
store.upsert_document("3", [make_node("c1", "3", seed=100.0)])
|
store.upsert_document(3, [make_node("c1", 3, seed=100.0)])
|
||||||
assert "c1" in _query(store, [100.0] * DIM, top_k=1).ids
|
assert "c1" in _query(store, [100.0] * DIM, top_k=1).ids
|
||||||
|
|
||||||
def test_auto_compact_triggers_on_churn(self, store) -> None:
|
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)])
|
store.add([make_node(f"s{j}", 1, seed=float(j)) for j in range(20)])
|
||||||
self._churn(store, 5)
|
self._churn(store, 5)
|
||||||
assert self._bloat_ratio(store) > 2
|
assert self._bloat_ratio(store) > 2
|
||||||
store.compact()
|
store.compact()
|
||||||
@@ -393,7 +473,7 @@ class TestCompact:
|
|||||||
but a concurrent reader keeps -wal/-shm alive, so the cleanup must
|
but a concurrent reader keeps -wal/-shm alive, so the cleanup must
|
||||||
unlink them explicitly (as the structural-migration path does).
|
unlink them explicitly (as the structural-migration path does).
|
||||||
"""
|
"""
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
compact_path = str(tmp_path / DB_FILENAME) + ".compact"
|
compact_path = str(tmp_path / DB_FILENAME) + ".compact"
|
||||||
held: list[sqlite3.Connection] = []
|
held: list[sqlite3.Connection] = []
|
||||||
|
|
||||||
@@ -428,17 +508,76 @@ class TestCompact:
|
|||||||
A tiny batch size forces several fetchmany()/executemany() cycles so a
|
A tiny batch size forces several fetchmany()/executemany() cycles so a
|
||||||
regression in the streaming loop (dropped tail, off-by-one) surfaces.
|
regression in the streaming loop (dropped tail, off-by-one) surfaces.
|
||||||
"""
|
"""
|
||||||
monkeypatch.setattr("paperless_ai.vector_store.COMPACT_BATCH_SIZE", 3)
|
monkeypatch.setattr("paperless_ai.vector_store.BATCH_SIZE", 3)
|
||||||
store.add([make_node(f"n{i}", "1", seed=float(i)) for i in range(10)])
|
store.add([make_node(f"n{i}", 1, seed=float(i)) for i in range(10)])
|
||||||
store.compact(force=True)
|
store.compact(force=True)
|
||||||
ids = {n.node_id for n in store.get_nodes(filters=_in_filter(["1"]))}
|
ids = {n.node_id for n in store.get_nodes(filters=_in_filter([1]))}
|
||||||
assert ids == {f"n{i}" for i in range(10)}
|
assert ids == {f"n{i}" for i in range(10)}
|
||||||
assert self._bloat_ratio(store) == pytest.approx(1.0)
|
assert self._bloat_ratio(store) == pytest.approx(1.0)
|
||||||
|
|
||||||
|
def test_force_compact_preserves_modified_times(
|
||||||
|
self,
|
||||||
|
store: PaperlessSqliteVecVectorStore,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- A store with documents whose modified times are tracked
|
||||||
|
WHEN:
|
||||||
|
- compact(force=True) rebuilds the database file
|
||||||
|
THEN:
|
||||||
|
- get_modified_times() still returns every document's value
|
||||||
|
(document_meta must be copied across the file-swap, not just
|
||||||
|
the vec0 rows)
|
||||||
|
"""
|
||||||
|
store.add(
|
||||||
|
[
|
||||||
|
make_node("a1", 1, modified="2026-01-01T00:00:00"),
|
||||||
|
make_node("b1", 2, modified="2026-02-02T00:00:00"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
before = store.get_modified_times()
|
||||||
|
store.compact(force=True)
|
||||||
|
assert store.get_modified_times() == before
|
||||||
|
|
||||||
|
def test_compact_on_unmigrated_store_is_noop(
|
||||||
|
self,
|
||||||
|
store: PaperlessSqliteVecVectorStore,
|
||||||
|
mocker: MockerFixture,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- A store whose schema_version has been forced behind
|
||||||
|
SCHEMA_VERSION (has_pending_migration() is True)
|
||||||
|
WHEN:
|
||||||
|
- compact(force=True) is called directly, without the caller
|
||||||
|
having run check_and_run_migrations() first
|
||||||
|
THEN:
|
||||||
|
- compact() is a safe no-op: no file-swap rebuild is attempted
|
||||||
|
at all (asserted via a spy on _rebuild_into, since schema_version
|
||||||
|
alone is not a reliable signal -- a rebuild would otherwise
|
||||||
|
copy the stale schema_version across unchanged, making a
|
||||||
|
before/after equality check pass even when a rebuild *did*
|
||||||
|
happen). Rebuilding an unmigrated store would silently lose
|
||||||
|
document_meta and leave the swapped-in file claiming the old
|
||||||
|
schema_version -- see the class docstring rationale.
|
||||||
|
"""
|
||||||
|
store.add([make_node("a1", 1)])
|
||||||
|
store.client.execute(
|
||||||
|
"UPDATE index_meta SET value = '0' WHERE key = 'schema_version'",
|
||||||
|
)
|
||||||
|
assert store.has_pending_migration() is True
|
||||||
|
rebuild_spy = mocker.spy(PaperlessSqliteVecVectorStore, "_rebuild_into")
|
||||||
|
store.compact(force=True)
|
||||||
|
rebuild_spy.assert_not_called()
|
||||||
|
row = store.client.execute(
|
||||||
|
"SELECT value FROM index_meta WHERE key = 'schema_version'",
|
||||||
|
).fetchone()
|
||||||
|
assert int(row["value"]) == 0
|
||||||
|
|
||||||
|
|
||||||
class TestDbFile:
|
class TestDbFile:
|
||||||
def test_single_db_file_in_index_dir(self, store, tmp_path: Path) -> None:
|
def test_single_db_file_in_index_dir(self, store, tmp_path: Path) -> None:
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
assert (tmp_path / DB_FILENAME).exists()
|
assert (tmp_path / DB_FILENAME).exists()
|
||||||
|
|
||||||
def test_wal_mode_enabled(self, store) -> None:
|
def test_wal_mode_enabled(self, store) -> None:
|
||||||
@@ -448,7 +587,26 @@ class TestDbFile:
|
|||||||
|
|
||||||
|
|
||||||
class TestMigrations:
|
class TestMigrations:
|
||||||
"""Tests for the schema migration machinery."""
|
"""Tests for the schema migration machinery.
|
||||||
|
|
||||||
|
These tests exercise check_and_run_migrations()'s generic dispatch logic
|
||||||
|
(structural vs. re-embed, version-boundary stopping) using ad hoc test
|
||||||
|
migrations layered on top of SCHEMA_VERSION -- distinct from
|
||||||
|
TestV1ToV2Migration, which exercises the real, frozen m0001_v1_to_v2
|
||||||
|
migration. Test migrations use version numbers starting at
|
||||||
|
SCHEMA_VERSION (2) and above so they never collide with the real
|
||||||
|
from_version=1/to_version=2 migration already registered in MIGRATIONS.
|
||||||
|
|
||||||
|
The fake structural migrations' apply() fixtures (see
|
||||||
|
test_structural_migration_copies_rows_and_updates_version and
|
||||||
|
test_stop_at_reembed_boundary below) only populate the rebuilt vec0
|
||||||
|
table itself -- they never insert into document_chunks/document_meta on
|
||||||
|
the destination connection. That's fine here: these tests exist to
|
||||||
|
verify the generic dispatch mechanism (version bookkeeping, structural-
|
||||||
|
vs-reembed branching), not full schema correctness of a rebuilt store;
|
||||||
|
the real migration's data completeness is covered separately by
|
||||||
|
TestV1ToV2Migration.
|
||||||
|
"""
|
||||||
|
|
||||||
def _schema_version(self, store: PaperlessSqliteVecVectorStore) -> int | None:
|
def _schema_version(self, store: PaperlessSqliteVecVectorStore) -> int | None:
|
||||||
row = store.client.execute(
|
row = store.client.execute(
|
||||||
@@ -457,21 +615,21 @@ class TestMigrations:
|
|||||||
return int(row[0]) if row else None
|
return int(row[0]) if row else None
|
||||||
|
|
||||||
def test_new_table_records_schema_version(self, store) -> None:
|
def test_new_table_records_schema_version(self, store) -> None:
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
assert self._schema_version(store) == SCHEMA_VERSION
|
assert self._schema_version(store) == SCHEMA_VERSION
|
||||||
|
|
||||||
def test_check_migrations_no_table_returns_false(self, store) -> None:
|
def test_check_migrations_no_table_returns_false(self, store) -> None:
|
||||||
assert store.check_and_run_migrations() is False
|
assert store.check_and_run_migrations() is False
|
||||||
|
|
||||||
def test_check_migrations_current_version_returns_false(self, store) -> None:
|
def test_check_migrations_current_version_returns_false(self, store) -> None:
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
assert store.check_and_run_migrations() is False
|
assert store.check_and_run_migrations() is False
|
||||||
|
|
||||||
def test_reembed_migration_returns_true(self, store, tmp_path: Path) -> None:
|
def test_reembed_migration_returns_true(self, store, tmp_path: Path) -> None:
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
migration = Migration(
|
migration = Migration(
|
||||||
from_version=1,
|
from_version=SCHEMA_VERSION,
|
||||||
to_version=2,
|
to_version=SCHEMA_VERSION + 1,
|
||||||
kind="re-embed",
|
kind="re-embed",
|
||||||
description="test re-embed",
|
description="test re-embed",
|
||||||
)
|
)
|
||||||
@@ -480,7 +638,7 @@ class TestMigrations:
|
|||||||
from paperless_ai import vector_store as vs_mod
|
from paperless_ai import vector_store as vs_mod
|
||||||
|
|
||||||
original = vs_mod.SCHEMA_VERSION
|
original = vs_mod.SCHEMA_VERSION
|
||||||
vs_mod.SCHEMA_VERSION = 2
|
vs_mod.SCHEMA_VERSION = SCHEMA_VERSION + 1
|
||||||
result = store.check_and_run_migrations()
|
result = store.check_and_run_migrations()
|
||||||
finally:
|
finally:
|
||||||
MIGRATIONS.remove(migration)
|
MIGRATIONS.remove(migration)
|
||||||
@@ -492,7 +650,7 @@ class TestMigrations:
|
|||||||
store,
|
store,
|
||||||
tmp_path: Path,
|
tmp_path: Path,
|
||||||
) -> None:
|
) -> None:
|
||||||
store.add([make_node("a1", "1"), make_node("b1", "2")])
|
store.add([make_node("a1", 1), make_node("b1", 2)])
|
||||||
|
|
||||||
def apply(
|
def apply(
|
||||||
src: sqlite3.Connection,
|
src: sqlite3.Connection,
|
||||||
@@ -511,7 +669,7 @@ class TestMigrations:
|
|||||||
(str(dim),),
|
(str(dim),),
|
||||||
)
|
)
|
||||||
rows = src.execute(
|
rows = src.execute(
|
||||||
"SELECT id, document_id, modified, node_content, embedding "
|
"SELECT id, document_id, node_content, embedding "
|
||||||
f"FROM {DEFAULT_TABLE_NAME}",
|
f"FROM {DEFAULT_TABLE_NAME}",
|
||||||
).fetchall()
|
).fetchall()
|
||||||
dst.execute("BEGIN IMMEDIATE")
|
dst.execute("BEGIN IMMEDIATE")
|
||||||
@@ -522,8 +680,8 @@ class TestMigrations:
|
|||||||
[
|
[
|
||||||
(
|
(
|
||||||
r["id"],
|
r["id"],
|
||||||
r["document_id"],
|
str(r["document_id"]),
|
||||||
r["modified"],
|
"",
|
||||||
r["node_content"],
|
r["node_content"],
|
||||||
bytes(r["embedding"]),
|
bytes(r["embedding"]),
|
||||||
)
|
)
|
||||||
@@ -538,8 +696,8 @@ class TestMigrations:
|
|||||||
dst.execute("COMMIT")
|
dst.execute("COMMIT")
|
||||||
|
|
||||||
migration = Migration(
|
migration = Migration(
|
||||||
from_version=1,
|
from_version=SCHEMA_VERSION,
|
||||||
to_version=2,
|
to_version=SCHEMA_VERSION + 1,
|
||||||
kind="structural",
|
kind="structural",
|
||||||
description="test structural",
|
description="test structural",
|
||||||
apply=apply,
|
apply=apply,
|
||||||
@@ -549,28 +707,29 @@ class TestMigrations:
|
|||||||
from paperless_ai import vector_store as vs_mod
|
from paperless_ai import vector_store as vs_mod
|
||||||
|
|
||||||
original = vs_mod.SCHEMA_VERSION
|
original = vs_mod.SCHEMA_VERSION
|
||||||
vs_mod.SCHEMA_VERSION = 2
|
vs_mod.SCHEMA_VERSION = SCHEMA_VERSION + 1
|
||||||
result = store.check_and_run_migrations()
|
result = store.check_and_run_migrations()
|
||||||
finally:
|
finally:
|
||||||
MIGRATIONS.remove(migration)
|
MIGRATIONS.remove(migration)
|
||||||
vs_mod.SCHEMA_VERSION = original
|
vs_mod.SCHEMA_VERSION = original
|
||||||
|
|
||||||
assert result is False
|
assert result is False
|
||||||
assert self._schema_version(store) == 2
|
assert self._schema_version(store) == SCHEMA_VERSION + 1
|
||||||
ids = {n.node_id for n in store.get_nodes()}
|
ids = {n.node_id for n in store.get_nodes()}
|
||||||
assert ids == {"a1", "b1"}
|
assert ids == {"a1", "b1"}
|
||||||
|
|
||||||
def test_compact_preserves_schema_version(self, store) -> None:
|
def test_compact_preserves_schema_version(self, store) -> None:
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
assert self._schema_version(store) == SCHEMA_VERSION
|
assert self._schema_version(store) == SCHEMA_VERSION
|
||||||
store.compact(force=True)
|
store.compact(force=True)
|
||||||
assert self._schema_version(store) == SCHEMA_VERSION
|
assert self._schema_version(store) == SCHEMA_VERSION
|
||||||
|
|
||||||
def test_stop_at_reembed_boundary(self, store) -> None:
|
def test_stop_at_reembed_boundary(self, store) -> None:
|
||||||
# Registry: structural v2, re-embed v3, structural v4.
|
# Registry: structural v(N+1), re-embed v(N+2), structural v(N+3),
|
||||||
# Only v2 should apply; the re-embed boundary must stop execution
|
# where N = SCHEMA_VERSION. Only v(N+1) should apply; the re-embed
|
||||||
# before v4 runs, and the stored version must stay at 2.
|
# boundary must stop execution before v(N+3) runs, and the stored
|
||||||
store.add([make_node("a1", "1"), make_node("b1", "2")])
|
# version must stay at N+1.
|
||||||
|
store.add([make_node("a1", 1), make_node("b1", 2)])
|
||||||
|
|
||||||
def copy_apply(
|
def copy_apply(
|
||||||
src: sqlite3.Connection,
|
src: sqlite3.Connection,
|
||||||
@@ -589,7 +748,7 @@ class TestMigrations:
|
|||||||
(str(dim),),
|
(str(dim),),
|
||||||
)
|
)
|
||||||
rows = src.execute(
|
rows = src.execute(
|
||||||
"SELECT id, document_id, modified, node_content, embedding "
|
"SELECT id, document_id, node_content, embedding "
|
||||||
f"FROM {DEFAULT_TABLE_NAME}",
|
f"FROM {DEFAULT_TABLE_NAME}",
|
||||||
).fetchall()
|
).fetchall()
|
||||||
dst.execute("BEGIN IMMEDIATE")
|
dst.execute("BEGIN IMMEDIATE")
|
||||||
@@ -600,8 +759,8 @@ class TestMigrations:
|
|||||||
[
|
[
|
||||||
(
|
(
|
||||||
r["id"],
|
r["id"],
|
||||||
r["document_id"],
|
str(r["document_id"]),
|
||||||
r["modified"],
|
"",
|
||||||
r["node_content"],
|
r["node_content"],
|
||||||
bytes(r["embedding"]),
|
bytes(r["embedding"]),
|
||||||
)
|
)
|
||||||
@@ -612,23 +771,23 @@ class TestMigrations:
|
|||||||
|
|
||||||
migrations = [
|
migrations = [
|
||||||
Migration(
|
Migration(
|
||||||
from_version=1,
|
from_version=SCHEMA_VERSION,
|
||||||
to_version=2,
|
to_version=SCHEMA_VERSION + 1,
|
||||||
kind="structural",
|
kind="structural",
|
||||||
description="v2 structural",
|
description="v(N+1) structural",
|
||||||
apply=copy_apply,
|
apply=copy_apply,
|
||||||
),
|
),
|
||||||
Migration(
|
Migration(
|
||||||
from_version=2,
|
from_version=SCHEMA_VERSION + 1,
|
||||||
to_version=3,
|
to_version=SCHEMA_VERSION + 2,
|
||||||
kind="re-embed",
|
kind="re-embed",
|
||||||
description="v3 re-embed boundary",
|
description="v(N+2) re-embed boundary",
|
||||||
),
|
),
|
||||||
Migration(
|
Migration(
|
||||||
from_version=3,
|
from_version=SCHEMA_VERSION + 2,
|
||||||
to_version=4,
|
to_version=SCHEMA_VERSION + 3,
|
||||||
kind="structural",
|
kind="structural",
|
||||||
description="v4 structural - must not run",
|
description="v(N+3) structural - must not run",
|
||||||
apply=copy_apply,
|
apply=copy_apply,
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
@@ -637,7 +796,7 @@ class TestMigrations:
|
|||||||
from paperless_ai import vector_store as vs_mod
|
from paperless_ai import vector_store as vs_mod
|
||||||
|
|
||||||
original = vs_mod.SCHEMA_VERSION
|
original = vs_mod.SCHEMA_VERSION
|
||||||
vs_mod.SCHEMA_VERSION = 4
|
vs_mod.SCHEMA_VERSION = SCHEMA_VERSION + 3
|
||||||
result = store.check_and_run_migrations()
|
result = store.check_and_run_migrations()
|
||||||
finally:
|
finally:
|
||||||
for m in migrations:
|
for m in migrations:
|
||||||
@@ -645,7 +804,7 @@ class TestMigrations:
|
|||||||
vs_mod.SCHEMA_VERSION = original
|
vs_mod.SCHEMA_VERSION = original
|
||||||
|
|
||||||
assert result is True
|
assert result is True
|
||||||
assert self._schema_version(store) == 2
|
assert self._schema_version(store) == SCHEMA_VERSION + 1
|
||||||
|
|
||||||
def test_has_pending_migration_false_when_no_table(
|
def test_has_pending_migration_false_when_no_table(
|
||||||
self,
|
self,
|
||||||
@@ -673,7 +832,7 @@ class TestMigrations:
|
|||||||
THEN:
|
THEN:
|
||||||
- False is returned
|
- False is returned
|
||||||
"""
|
"""
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
assert store.has_pending_migration() is False
|
assert store.has_pending_migration() is False
|
||||||
|
|
||||||
def test_has_pending_migration_true_when_behind(
|
def test_has_pending_migration_true_when_behind(
|
||||||
@@ -688,8 +847,183 @@ class TestMigrations:
|
|||||||
THEN:
|
THEN:
|
||||||
- True is returned
|
- True is returned
|
||||||
"""
|
"""
|
||||||
store.add([make_node("a1", "1")])
|
store.add([make_node("a1", 1)])
|
||||||
store.client.execute(
|
store.client.execute(
|
||||||
"UPDATE index_meta SET value = '0' WHERE key = 'schema_version'",
|
"UPDATE index_meta SET value = '0' WHERE key = 'schema_version'",
|
||||||
)
|
)
|
||||||
assert store.has_pending_migration() is True
|
assert store.has_pending_migration() is True
|
||||||
|
|
||||||
|
|
||||||
|
class TestV1ToV2Migration:
|
||||||
|
"""m0001_v1_to_v2 migrates a real, historically-shaped v1 store. The
|
||||||
|
fixture below is a literal, hardcoded v1 DDL string -- NOT derived from
|
||||||
|
any current code -- so this test keeps testing the actual historical
|
||||||
|
shape even if vector_store.py's "current" schema changes again later.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def _build_v1_store(self, db_path: str, dim: int) -> None:
|
||||||
|
conn = sqlite3.connect(db_path)
|
||||||
|
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)",
|
||||||
|
)
|
||||||
|
conn.execute( # nosemgrep
|
||||||
|
"CREATE VIRTUAL TABLE documents USING vec0("
|
||||||
|
"id TEXT PRIMARY KEY, document_id TEXT, modified TEXT,"
|
||||||
|
f" +node_content TEXT, embedding float[{dim}] distance_metric=cosine"
|
||||||
|
")",
|
||||||
|
)
|
||||||
|
conn.execute(
|
||||||
|
"INSERT INTO index_meta (key, value) VALUES ('dim', ?)",
|
||||||
|
(str(dim),),
|
||||||
|
)
|
||||||
|
conn.execute(
|
||||||
|
"INSERT INTO index_meta (key, value) VALUES ('schema_version', '1')",
|
||||||
|
)
|
||||||
|
conn.execute(
|
||||||
|
"INSERT INTO index_meta (key, value) VALUES ('embed_model', 'model-a')",
|
||||||
|
)
|
||||||
|
rows = [
|
||||||
|
("c1", "1", "2026-01-01T00:00:00", '{"text": "a"}', _pack([0.1] * dim)),
|
||||||
|
("c2", "1", "2026-01-01T00:00:00", '{"text": "b"}', _pack([0.2] * dim)),
|
||||||
|
("c3", "2", "2026-02-02T00:00:00", '{"text": "c"}', _pack([0.3] * dim)),
|
||||||
|
]
|
||||||
|
conn.executemany(
|
||||||
|
"INSERT INTO documents (id, document_id, modified, node_content, embedding)"
|
||||||
|
" VALUES (?, ?, ?, ?, ?)",
|
||||||
|
rows,
|
||||||
|
)
|
||||||
|
conn.execute(
|
||||||
|
"INSERT INTO index_meta (key, value) VALUES ('total_inserts', '3')",
|
||||||
|
)
|
||||||
|
conn.commit()
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
def test_migration_converts_v1_store_to_v2(self, tmp_path: Path) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- A real v1-shaped store (TEXT document_id, modified inline in
|
||||||
|
vec0, no document_chunks/document_meta) built from a literal,
|
||||||
|
hardcoded historical DDL
|
||||||
|
WHEN:
|
||||||
|
- A PaperlessSqliteVecVectorStore is opened against it
|
||||||
|
THEN:
|
||||||
|
- schema_version becomes 2, document_id values become int,
|
||||||
|
document_chunks/document_meta are backfilled once per chunk/
|
||||||
|
document respectively, and dim/embed_model survive
|
||||||
|
"""
|
||||||
|
db_dir = tmp_path
|
||||||
|
self._build_v1_store(str(db_dir / DB_FILENAME), dim=16)
|
||||||
|
with PaperlessSqliteVecVectorStore(uri=str(db_dir)) as store:
|
||||||
|
assert store.check_and_run_migrations() is False
|
||||||
|
row = store.client.execute(
|
||||||
|
"SELECT value FROM index_meta WHERE key = 'schema_version'",
|
||||||
|
).fetchone()
|
||||||
|
assert int(row["value"]) == 2
|
||||||
|
doc_id_row = store.client.execute(
|
||||||
|
"SELECT document_id FROM documents WHERE id = 'c1'",
|
||||||
|
).fetchone()
|
||||||
|
assert isinstance(doc_id_row["document_id"], int)
|
||||||
|
assert doc_id_row["document_id"] == 1
|
||||||
|
chunk_ids = sorted(
|
||||||
|
r["chunk_id"]
|
||||||
|
for r in store.client.execute(
|
||||||
|
"SELECT chunk_id FROM document_chunks",
|
||||||
|
).fetchall()
|
||||||
|
)
|
||||||
|
assert chunk_ids == ["c1", "c2", "c3"]
|
||||||
|
assert store.get_modified_times() == {
|
||||||
|
"1": "2026-01-01T00:00:00",
|
||||||
|
"2": "2026-02-02T00:00:00",
|
||||||
|
}
|
||||||
|
assert store.stored_model_name() == "model-a"
|
||||||
|
assert store.vector_dim() == 16
|
||||||
|
|
||||||
|
def test_migration_raises_on_malformed_document_id(
|
||||||
|
self,
|
||||||
|
tmp_path: Path,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- A v1-shaped store with a corrupted, non-integer document_id
|
||||||
|
value on one row
|
||||||
|
WHEN:
|
||||||
|
- The migration runs
|
||||||
|
THEN:
|
||||||
|
- A ValueError is raised (fail loudly, no silent data loss) --
|
||||||
|
this matches the rest of vector_store.py, which has no
|
||||||
|
precedent for silently skipping malformed rows
|
||||||
|
"""
|
||||||
|
db_dir = tmp_path
|
||||||
|
self._build_v1_store(str(db_dir / DB_FILENAME), dim=16)
|
||||||
|
conn = sqlite3.connect(str(db_dir / DB_FILENAME))
|
||||||
|
conn.enable_load_extension(True) # noqa: FBT003
|
||||||
|
sqlite_vec.load(conn)
|
||||||
|
conn.enable_load_extension(False) # noqa: FBT003
|
||||||
|
conn.execute(
|
||||||
|
"UPDATE documents SET document_id = 'not-an-int' WHERE id = 'c1'",
|
||||||
|
)
|
||||||
|
conn.commit()
|
||||||
|
conn.close()
|
||||||
|
with (
|
||||||
|
pytest.raises(ValueError),
|
||||||
|
PaperlessSqliteVecVectorStore(uri=str(db_dir)) as store,
|
||||||
|
):
|
||||||
|
store.check_and_run_migrations()
|
||||||
|
|
||||||
|
def test_migration_never_delegates_to_current_schema_helpers(
|
||||||
|
self,
|
||||||
|
tmp_path: Path,
|
||||||
|
mocker: MockerFixture,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- A real v1-shaped store
|
||||||
|
WHEN:
|
||||||
|
- The migration runs, with DocumentChunksTable.create/
|
||||||
|
DocumentMetaTable.create/_create_vec_table spied on
|
||||||
|
THEN:
|
||||||
|
- None of those "current schema" helpers are ever called during
|
||||||
|
the migration -- it must freeze its own historical DDL, per
|
||||||
|
the DDL-freezing rule (see spec), so a future schema bump
|
||||||
|
can't silently corrupt this migration's output
|
||||||
|
"""
|
||||||
|
db_dir = tmp_path
|
||||||
|
self._build_v1_store(str(db_dir / DB_FILENAME), dim=16)
|
||||||
|
chunks_create_spy = mocker.spy(DocumentChunksTable, "create")
|
||||||
|
meta_create_spy = mocker.spy(DocumentMetaTable, "create")
|
||||||
|
vec_table_spy = mocker.spy(
|
||||||
|
PaperlessSqliteVecVectorStore,
|
||||||
|
"_create_vec_table",
|
||||||
|
)
|
||||||
|
with PaperlessSqliteVecVectorStore(uri=str(db_dir)) as store:
|
||||||
|
store.check_and_run_migrations()
|
||||||
|
# _open_connection() legitimately calls create() three times across
|
||||||
|
# a structural migration: once for the store's own live connection
|
||||||
|
# (construction), once for the migration's temp rebuild file
|
||||||
|
# (_rebuild_file), and once more when _swap_in_compact() reopens the
|
||||||
|
# swapped-in file as self._conn. What matters is that
|
||||||
|
# m0001_v1_to_v2's apply() itself never calls these directly, which
|
||||||
|
# a source-text check alone can't prove (it can't tell "mentioned in
|
||||||
|
# a comment" from "actually called", and is trivially evadable by
|
||||||
|
# importing a symbol under an alias). Asserting the exact call count
|
||||||
|
# instead: exactly 3 calls to each create() (all from
|
||||||
|
# _open_connection, never a 4th from inside apply()), and zero calls
|
||||||
|
# to _create_vec_table (neither _open_connection nor apply() calls
|
||||||
|
# it -- apply() freezes its own literal CREATE VIRTUAL TABLE DDL
|
||||||
|
# instead).
|
||||||
|
assert chunks_create_spy.call_count == 3
|
||||||
|
assert meta_create_spy.call_count == 3
|
||||||
|
assert vec_table_spy.call_count == 0
|
||||||
|
# Cheap secondary signal, kept alongside the spy assertions above
|
||||||
|
# (not in place of them): the migration module's source should never
|
||||||
|
# even mention these "current schema" helpers by name.
|
||||||
|
source = inspect.getsource(m0001_v1_to_v2)
|
||||||
|
assert "DocumentChunksTable.create" not in source
|
||||||
|
assert "DocumentMetaTable.create" not in source
|
||||||
|
assert "_create_vec_table(" not in source
|
||||||
|
|||||||
+276
-158
@@ -8,6 +8,7 @@ from contextlib import contextmanager
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from types import TracebackType
|
from types import TracebackType
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
from typing import NamedTuple
|
||||||
|
|
||||||
import sqlite_vec
|
import sqlite_vec
|
||||||
from llama_index.core.bridge.pydantic import PrivateAttr
|
from llama_index.core.bridge.pydantic import PrivateAttr
|
||||||
@@ -24,6 +25,11 @@ from llama_index.core.vector_stores.utils import node_to_metadata_dict
|
|||||||
|
|
||||||
from paperless_ai.migrations import MIGRATIONS
|
from paperless_ai.migrations import MIGRATIONS
|
||||||
from paperless_ai.migrations import Migration
|
from paperless_ai.migrations import Migration
|
||||||
|
from paperless_ai.tables import ChunkRow
|
||||||
|
from paperless_ai.tables import DocumentChunksTable
|
||||||
|
from paperless_ai.tables import DocumentMetaRow
|
||||||
|
from paperless_ai.tables import DocumentMetaTable
|
||||||
|
from paperless_ai.tables import IndexMetaTable
|
||||||
|
|
||||||
logger = logging.getLogger("paperless_ai.vector_store")
|
logger = logging.getLogger("paperless_ai.vector_store")
|
||||||
|
|
||||||
@@ -33,7 +39,7 @@ DEFAULT_TABLE_NAME = "documents"
|
|||||||
# Current schema version. Written to index_meta at table creation and bumped
|
# Current schema version. Written to index_meta at table creation and bumped
|
||||||
# whenever a Migration is added to MIGRATIONS. check_and_run_migrations() uses
|
# whenever a Migration is added to MIGRATIONS. check_and_run_migrations() uses
|
||||||
# this to decide which migrations to run on an existing store.
|
# this to decide which migrations to run on an existing store.
|
||||||
SCHEMA_VERSION = 1
|
SCHEMA_VERSION = 2
|
||||||
|
|
||||||
# compact(): rebuild when the cumulative rowid count exceeds this multiple of
|
# compact(): rebuild when the cumulative rowid count exceeds this multiple of
|
||||||
# the live row count. DELETEs on vec0 tables never reclaim space (upstream
|
# the live row count. DELETEs on vec0 tables never reclaim space (upstream
|
||||||
@@ -41,15 +47,32 @@ SCHEMA_VERSION = 1
|
|||||||
# a rebuild copies the live rows into a fresh table.
|
# a rebuild copies the live rows into a fresh table.
|
||||||
COMPACT_BLOAT_RATIO = 2.0
|
COMPACT_BLOAT_RATIO = 2.0
|
||||||
|
|
||||||
# compact(): number of rows copied per executemany() when rebuilding the file.
|
# Number of rows fetched/copied per batch whenever this module streams rows
|
||||||
# Rows are streamed from the source cursor in batches of this size rather than
|
# instead of materializing them all at once, keeping memory bounded regardless
|
||||||
# materialized all at once, keeping memory bounded regardless of index size.
|
# of index size -- used by compact()'s rebuild, m0001_v1_to_v2's migration
|
||||||
COMPACT_BATCH_SIZE = 500
|
# copy, and DocumentMetaTable.copy_all(). No longer compact()-specific, hence
|
||||||
|
# the plain name.
|
||||||
|
BATCH_SIZE = 500
|
||||||
|
|
||||||
# Filterable vec0 metadata columns. _build_where() only ever receives filter
|
# Filterable vec0 metadata columns. _build_where() only ever receives filter
|
||||||
# keys we construct ourselves, but allowlisting keeps SQL identifiers safe by
|
# keys we construct ourselves, but allowlisting keeps SQL identifiers safe by
|
||||||
# construction.
|
# construction. "modified" is not here: it is never filtered on, and as of
|
||||||
_FILTER_COLUMNS = frozenset({"document_id", "modified"})
|
# schema v2 it isn't even a vec0 column anymore (see document_meta).
|
||||||
|
_FILTER_COLUMNS = frozenset({"document_id"})
|
||||||
|
|
||||||
|
|
||||||
|
class _Row(NamedTuple):
|
||||||
|
"""One node, ready to write. ``modified`` is not a vec0 column (see
|
||||||
|
document_meta) -- it rides along here because every row-producing call
|
||||||
|
site needs both the vec0 insert values and the document_meta upsert
|
||||||
|
value from the same node.
|
||||||
|
"""
|
||||||
|
|
||||||
|
chunk_id: str
|
||||||
|
document_id: int
|
||||||
|
modified: str
|
||||||
|
node_content: str
|
||||||
|
embedding: bytes
|
||||||
|
|
||||||
|
|
||||||
def _pack(embedding: Sequence[float]) -> bytes:
|
def _pack(embedding: Sequence[float]) -> bytes:
|
||||||
@@ -60,14 +83,30 @@ def _unpack(blob: bytes) -> list[float]:
|
|||||||
return list(struct.unpack(f"{len(blob) // 4}f", blob))
|
return list(struct.unpack(f"{len(blob) // 4}f", blob))
|
||||||
|
|
||||||
|
|
||||||
def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]:
|
_INSERT = (
|
||||||
"""Translate the EQ / IN / NE filters we use into a parameterized SQL clause
|
"INSERT INTO "
|
||||||
on vec0 metadata columns. Returns ("", []) when there is nothing to filter.
|
+ DEFAULT_TABLE_NAME
|
||||||
|
+ " (id, document_id, node_content, embedding) VALUES (?, ?, ?, ?)"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _vec0_params(rows: list[_Row]) -> list[tuple[str, int, str, bytes]]:
|
||||||
|
"""``rows``, minus the ``modified`` field vec0 no longer stores."""
|
||||||
|
return [(r.chunk_id, r.document_id, r.node_content, r.embedding) for r in rows]
|
||||||
|
|
||||||
|
|
||||||
|
def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]:
|
||||||
|
"""Translate the EQ / IN / NE filters we use into a parameterized SQL
|
||||||
|
clause on vec0 metadata columns. Returns ("", []) when there is nothing
|
||||||
|
to filter. document_id is vec0's only filterable column and is INTEGER;
|
||||||
|
every value is coerced via int() here so callers (which today still pass
|
||||||
|
strings in places, e.g. indexing.py's MetadataFilter construction) don't
|
||||||
|
have to be individually correct -- vec0 doesn't coerce types itself.
|
||||||
"""
|
"""
|
||||||
if filters is None or not filters.filters:
|
if filters is None or not filters.filters:
|
||||||
return "", []
|
return "", []
|
||||||
clauses: list[str] = []
|
clauses: list[str] = []
|
||||||
params: list[str] = []
|
params: list[int] = []
|
||||||
for f in filters.filters:
|
for f in filters.filters:
|
||||||
# filters.filters is Union[MetadataFilter, ExactMatchFilter, MetadataFilters];
|
# filters.filters is Union[MetadataFilter, ExactMatchFilter, MetadataFilters];
|
||||||
# we only build MetadataFilter entries, so skip anything else at runtime.
|
# we only build MetadataFilter entries, so skip anything else at runtime.
|
||||||
@@ -76,7 +115,7 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]:
|
|||||||
if f.key not in _FILTER_COLUMNS: # pragma: no cover - we build the keys
|
if f.key not in _FILTER_COLUMNS: # pragma: no cover - we build the keys
|
||||||
raise NotImplementedError(f"Unsupported filter column: {f.key}")
|
raise NotImplementedError(f"Unsupported filter column: {f.key}")
|
||||||
if f.operator == FilterOperator.IN:
|
if f.operator == FilterOperator.IN:
|
||||||
values = [str(v) for v in f.value] # type: ignore[union-attr] # value is list when operator is IN
|
values = [int(v) for v in f.value] # type: ignore[union-attr]
|
||||||
if not values: # pragma: no cover
|
if not values: # pragma: no cover
|
||||||
clauses.append("1 = 0")
|
clauses.append("1 = 0")
|
||||||
continue
|
continue
|
||||||
@@ -85,10 +124,10 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]:
|
|||||||
params.extend(values)
|
params.extend(values)
|
||||||
elif f.operator == FilterOperator.EQ:
|
elif f.operator == FilterOperator.EQ:
|
||||||
clauses.append(f"{f.key} = ?")
|
clauses.append(f"{f.key} = ?")
|
||||||
params.append(str(f.value))
|
params.append(int(f.value))
|
||||||
elif f.operator == FilterOperator.NE:
|
elif f.operator == FilterOperator.NE:
|
||||||
clauses.append(f"{f.key} != ?")
|
clauses.append(f"{f.key} != ?")
|
||||||
params.append(str(f.value))
|
params.append(int(f.value))
|
||||||
else: # pragma: no cover - we only ever build EQ/IN/NE filters
|
else: # pragma: no cover - we only ever build EQ/IN/NE filters
|
||||||
raise NotImplementedError(f"Unsupported filter operator: {f.operator}")
|
raise NotImplementedError(f"Unsupported filter operator: {f.operator}")
|
||||||
if not clauses:
|
if not clauses:
|
||||||
@@ -153,9 +192,21 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
conn.enable_load_extension(False) # noqa: FBT003
|
conn.enable_load_extension(False) # noqa: FBT003
|
||||||
conn.execute("PRAGMA journal_mode=WAL")
|
conn.execute("PRAGMA journal_mode=WAL")
|
||||||
conn.execute("PRAGMA synchronous=NORMAL")
|
conn.execute("PRAGMA synchronous=NORMAL")
|
||||||
conn.execute(
|
IndexMetaTable.create(conn)
|
||||||
"CREATE TABLE IF NOT EXISTS index_meta (key TEXT PRIMARY KEY, value TEXT)",
|
# vec0 metadata columns only get an efficient lookup path inside a
|
||||||
)
|
# KNN (MATCH) query; a plain `WHERE document_id = ?` is a full table
|
||||||
|
# scan regardless of index size. This plain, indexed table is how
|
||||||
|
# delete()/upsert_document() find a document's chunk ids without
|
||||||
|
# that scan.
|
||||||
|
DocumentChunksTable.create(conn)
|
||||||
|
# modified used to be a vec0 metadata column, but vec0 only inlines
|
||||||
|
# TEXT metadata up to 12 bytes -- an ISO timestamp is always longer,
|
||||||
|
# so every read recompiled and stepped a fresh SQL statement per row.
|
||||||
|
# It was never filtered on inside a KNN query either, so it never
|
||||||
|
# needed to be a vec0 column at all. One row per document here (not
|
||||||
|
# per chunk, like document_chunks), since every chunk of a document
|
||||||
|
# shares the same modified value -- see get_modified_times().
|
||||||
|
DocumentMetaTable.create(conn)
|
||||||
return conn
|
return conn
|
||||||
|
|
||||||
@property
|
@property
|
||||||
@@ -190,24 +241,6 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
else:
|
else:
|
||||||
self._conn.execute("COMMIT")
|
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:
|
def table_exists(self) -> bool:
|
||||||
return (
|
return (
|
||||||
self._conn.execute(
|
self._conn.execute(
|
||||||
@@ -220,18 +253,19 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
def vector_dim(self) -> int | None:
|
def vector_dim(self) -> int | None:
|
||||||
if not self.table_exists():
|
if not self.table_exists():
|
||||||
return None
|
return None
|
||||||
value = self._meta_get("dim")
|
return IndexMetaTable.get_dim(self._conn)
|
||||||
return int(value) if value else None
|
|
||||||
|
|
||||||
def drop_table(self) -> None:
|
def drop_table(self) -> None:
|
||||||
self._conn.execute("DROP TABLE IF EXISTS " + DEFAULT_TABLE_NAME)
|
self._conn.execute("DROP TABLE IF EXISTS " + DEFAULT_TABLE_NAME)
|
||||||
self._conn.execute("DELETE FROM index_meta")
|
self._conn.execute("DELETE FROM index_meta")
|
||||||
|
DocumentChunksTable.delete_all(self._conn)
|
||||||
|
DocumentMetaTable.delete_all(self._conn)
|
||||||
|
|
||||||
def stored_model_name(self) -> str | None:
|
def stored_model_name(self) -> str | None:
|
||||||
"""Return the embedding model name recorded at table creation, or None."""
|
"""Return the embedding model name recorded at table creation, or None."""
|
||||||
if not self.table_exists():
|
if not self.table_exists():
|
||||||
return None
|
return None
|
||||||
return self._meta_get("embed_model")
|
return IndexMetaTable.get_embed_model(self._conn)
|
||||||
|
|
||||||
def config_mismatch(self, model_name: str) -> bool:
|
def config_mismatch(self, model_name: str) -> bool:
|
||||||
"""True when the stored model name differs from ``model_name``.
|
"""True when the stored model name differs from ``model_name``.
|
||||||
@@ -249,14 +283,17 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
# document_id is deliberately a metadata column, NOT a partition key:
|
# document_id is deliberately a metadata column, NOT a partition key:
|
||||||
# partition keys change KNN `k` to per-partition semantics under IN
|
# partition keys change KNN `k` to per-partition semantics under IN
|
||||||
# filters (asg017/sqlite-vec#142); metadata columns give a correct
|
# filters (asg017/sqlite-vec#142); metadata columns give a correct
|
||||||
# global top-k.
|
# global top-k. INTEGER (not TEXT, as in schema v1): EQ/NE/IN
|
||||||
|
# comparisons become a native i64 array compare instead of per-row
|
||||||
|
# strncmp against a 16-byte text view, and this drops the unused
|
||||||
|
# metadatatext shadow table TEXT columns carry. modified is not a
|
||||||
|
# column here at all as of v2 -- see document_meta.
|
||||||
conn.execute( # nosemgrep: python.sqlalchemy.security.sqlalchemy-execute-raw-query.sqlalchemy-execute-raw-query
|
conn.execute( # nosemgrep: python.sqlalchemy.security.sqlalchemy-execute-raw-query.sqlalchemy-execute-raw-query
|
||||||
"CREATE VIRTUAL TABLE "
|
"CREATE VIRTUAL TABLE "
|
||||||
+ DEFAULT_TABLE_NAME
|
+ DEFAULT_TABLE_NAME
|
||||||
+ " USING vec0("
|
+ " USING vec0("
|
||||||
+ "id TEXT PRIMARY KEY,"
|
+ "id TEXT PRIMARY KEY,"
|
||||||
+ " document_id TEXT,"
|
+ " document_id INTEGER,"
|
||||||
+ " modified TEXT,"
|
|
||||||
+ " +node_content TEXT,"
|
+ " +node_content TEXT,"
|
||||||
+ " embedding float["
|
+ " embedding float["
|
||||||
+ str(int(dim))
|
+ str(int(dim))
|
||||||
@@ -266,37 +303,78 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
|
|
||||||
def _create_table(self, dim: int) -> None:
|
def _create_table(self, dim: int) -> None:
|
||||||
self._create_vec_table(self._conn, dim)
|
self._create_vec_table(self._conn, dim)
|
||||||
self._meta_set("dim", str(dim))
|
IndexMetaTable.set_dim(self._conn, dim)
|
||||||
self._meta_set("schema_version", str(SCHEMA_VERSION))
|
IndexMetaTable.set_schema_version(self._conn, SCHEMA_VERSION)
|
||||||
if self._embed_model_name:
|
if self._embed_model_name:
|
||||||
self._meta_set("embed_model", self._embed_model_name)
|
IndexMetaTable.set_embed_model(self._conn, self._embed_model_name)
|
||||||
|
|
||||||
def _ensure_table(self, dim: int) -> None:
|
def _ensure_table(self, dim: int, *, table_exists: bool) -> None:
|
||||||
if not self.table_exists():
|
if not table_exists:
|
||||||
self._create_table(dim)
|
self._create_table(dim)
|
||||||
|
|
||||||
def _row(self, node: BaseNode) -> tuple[str, str, str, str, bytes]:
|
def _row(self, node: BaseNode) -> _Row:
|
||||||
meta = node_to_metadata_dict(
|
meta = node_to_metadata_dict(
|
||||||
node,
|
node,
|
||||||
remove_text=False,
|
remove_text=False,
|
||||||
flat_metadata=self.flat_metadata,
|
flat_metadata=self.flat_metadata,
|
||||||
)
|
)
|
||||||
# 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")
|
document_id = node.ref_doc_id or node.metadata.get("document_id")
|
||||||
return (
|
return _Row(
|
||||||
node.node_id,
|
chunk_id=node.node_id,
|
||||||
str(document_id or ""),
|
# document_id is required -- int(None) raises TypeError and
|
||||||
str(node.metadata.get("modified") or ""),
|
# int("not-a-number") raises ValueError, both intentional:
|
||||||
json.dumps(meta),
|
# fail loudly on a malformed/missing document_id rather than
|
||||||
_pack(node.get_embedding()),
|
# silently indexing a chunk with no owning document. modified,
|
||||||
|
# below, still uses the str(x or "") sentinel pattern because a
|
||||||
|
# missing modified value is legitimate (vec0 no longer even
|
||||||
|
# stores it -- see document_meta), whereas document_id must
|
||||||
|
# always be present.
|
||||||
|
document_id=int(document_id),
|
||||||
|
modified=str(node.metadata.get("modified") or ""),
|
||||||
|
node_content=json.dumps(meta),
|
||||||
|
embedding=_pack(node.get_embedding()),
|
||||||
)
|
)
|
||||||
|
|
||||||
_INSERT = (
|
def _index_chunks(self, rows: list[_Row]) -> None:
|
||||||
"INSERT INTO "
|
"""Record each row's (chunk_id, document_id) in document_chunks, and
|
||||||
+ DEFAULT_TABLE_NAME
|
each row's (document_id, modified) in document_meta -- deduped
|
||||||
+ " (id, document_id, modified, node_content, embedding) VALUES (?, ?, ?, ?, ?)"
|
within the batch, since every chunk of a document shares the same
|
||||||
)
|
modified value -- kept in lockstep with every insert into the vec0
|
||||||
|
table.
|
||||||
|
"""
|
||||||
|
DocumentChunksTable.insert_many(
|
||||||
|
self._conn,
|
||||||
|
(ChunkRow(r.chunk_id, r.document_id) for r in rows),
|
||||||
|
)
|
||||||
|
modified_by_document = {r.document_id: r.modified for r in rows}
|
||||||
|
DocumentMetaTable.upsert_many(
|
||||||
|
self._conn,
|
||||||
|
(
|
||||||
|
DocumentMetaRow(doc_id, mod)
|
||||||
|
for doc_id, mod in modified_by_document.items()
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
def _delete_chunks_by_document_id(self, document_id: int) -> None:
|
||||||
|
"""Delete all of a document's chunks via point-deletes on `id`.
|
||||||
|
|
||||||
|
vec0 has no efficient lookup on the document_id metadata column
|
||||||
|
outside a KNN query, so a plain `DELETE ... WHERE document_id = ?`
|
||||||
|
is a full table scan regardless of index size. Looking the chunk
|
||||||
|
ids up in document_chunks first (a real indexed lookup) and
|
||||||
|
deleting each by its `id` primary key instead turns that scan into
|
||||||
|
a handful of O(1) point deletes.
|
||||||
|
"""
|
||||||
|
chunk_ids = DocumentChunksTable.chunk_ids_for_document(
|
||||||
|
self._conn,
|
||||||
|
document_id,
|
||||||
|
)
|
||||||
|
self._conn.executemany(
|
||||||
|
"DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE id = ?",
|
||||||
|
[(chunk_id,) for chunk_id in chunk_ids],
|
||||||
|
)
|
||||||
|
DocumentChunksTable.delete_for_document(self._conn, document_id)
|
||||||
|
DocumentMetaTable.delete_for_document(self._conn, document_id)
|
||||||
|
|
||||||
def _increment_total_inserts(self, count: int) -> None:
|
def _increment_total_inserts(self, count: int) -> None:
|
||||||
"""Increment the cumulative insert counter stored in index_meta.
|
"""Increment the cumulative insert counter stored in index_meta.
|
||||||
@@ -306,48 +384,56 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
live_rows exceeds COMPACT_BLOAT_RATIO the table has accumulated
|
live_rows exceeds COMPACT_BLOAT_RATIO the table has accumulated
|
||||||
enough deleted-but-not-freed rows to warrant a rebuild.
|
enough deleted-but-not-freed rows to warrant a rebuild.
|
||||||
"""
|
"""
|
||||||
current = int(self._meta_get("total_inserts") or "0")
|
IndexMetaTable.increment_total_inserts(self._conn, count)
|
||||||
self._meta_set("total_inserts", str(current + count))
|
|
||||||
|
|
||||||
def add(self, nodes: Sequence[BaseNode], **add_kwargs: Any) -> list[str]:
|
def add(self, nodes: Sequence[BaseNode], **add_kwargs: Any) -> list[str]:
|
||||||
if not nodes:
|
if not nodes:
|
||||||
return []
|
return []
|
||||||
rows = [self._row(node) for node in nodes]
|
rows = [self._row(node) for node in nodes]
|
||||||
with self._transaction():
|
with self._transaction():
|
||||||
self._ensure_table(len(nodes[0].get_embedding()))
|
self._ensure_table(
|
||||||
self._conn.executemany(self._INSERT, rows)
|
len(nodes[0].get_embedding()),
|
||||||
|
table_exists=self.table_exists(),
|
||||||
|
)
|
||||||
|
self._conn.executemany(_INSERT, _vec0_params(rows))
|
||||||
|
self._index_chunks(rows)
|
||||||
self._increment_total_inserts(len(rows))
|
self._increment_total_inserts(len(rows))
|
||||||
return [node.node_id for node in nodes]
|
return [node.node_id for node in nodes]
|
||||||
|
|
||||||
def upsert_document(self, document_id: str, nodes: list[BaseNode]) -> list[str]:
|
def upsert_document(
|
||||||
|
self,
|
||||||
|
document_id: int | str,
|
||||||
|
nodes: list[BaseNode],
|
||||||
|
) -> list[str]:
|
||||||
"""Atomically replace all stored chunks of ``document_id`` with ``nodes``.
|
"""Atomically replace all stored chunks of ``document_id`` with ``nodes``.
|
||||||
|
|
||||||
One transaction deletes the document's existing rows and inserts the
|
One transaction deletes the document's existing rows and inserts the
|
||||||
new set (vec0's INSERT OR REPLACE is broken upstream, #259, so
|
new set (vec0's INSERT OR REPLACE is broken upstream, so delete+insert
|
||||||
delete+insert it is). WAL readers in other processes see either the
|
it is). WAL readers in other processes see either the old or the new
|
||||||
old or the new chunk set, never a partial state.
|
chunk set, never a partial state.
|
||||||
"""
|
"""
|
||||||
|
doc_id = int(document_id)
|
||||||
rows = [self._row(node) for node in nodes]
|
rows = [self._row(node) for node in nodes]
|
||||||
with self._transaction():
|
with self._transaction():
|
||||||
if nodes:
|
table_exists = self.table_exists()
|
||||||
self._ensure_table(len(nodes[0].get_embedding()))
|
if nodes and not table_exists:
|
||||||
if self.table_exists():
|
self._ensure_table(
|
||||||
self._conn.execute(
|
len(nodes[0].get_embedding()),
|
||||||
"DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE document_id = ?",
|
table_exists=False,
|
||||||
(str(document_id),),
|
|
||||||
)
|
)
|
||||||
|
table_exists = True
|
||||||
|
if table_exists:
|
||||||
|
self._delete_chunks_by_document_id(doc_id)
|
||||||
if rows:
|
if rows:
|
||||||
self._conn.executemany(self._INSERT, rows)
|
self._conn.executemany(_INSERT, _vec0_params(rows))
|
||||||
|
self._index_chunks(rows)
|
||||||
self._increment_total_inserts(len(rows))
|
self._increment_total_inserts(len(rows))
|
||||||
return [node.node_id for node in nodes]
|
return [node.node_id for node in nodes]
|
||||||
|
|
||||||
def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
|
def delete(self, ref_doc_id: int | str, **delete_kwargs: Any) -> None:
|
||||||
if self.table_exists():
|
if self.table_exists():
|
||||||
with self._transaction():
|
with self._transaction():
|
||||||
self._conn.execute(
|
self._delete_chunks_by_document_id(int(ref_doc_id))
|
||||||
"DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE document_id = ?",
|
|
||||||
(str(ref_doc_id),),
|
|
||||||
)
|
|
||||||
|
|
||||||
def _rows_to_nodes(self, rows: list[sqlite3.Row]) -> list[BaseNode]:
|
def _rows_to_nodes(self, rows: list[sqlite3.Row]) -> list[BaseNode]:
|
||||||
nodes: list[BaseNode] = []
|
nodes: list[BaseNode] = []
|
||||||
@@ -417,41 +503,66 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
def get_modified_times(self) -> dict[str, str]:
|
def get_modified_times(self) -> dict[str, str]:
|
||||||
"""Return {document_id: stored_modified_isoformat} for all indexed documents.
|
"""Return {document_id: stored_modified_isoformat} for all indexed documents.
|
||||||
|
|
||||||
All chunks of a document share the same ``modified`` value, so the
|
document_meta already has exactly one row per document (not per
|
||||||
first row seen per document is sufficient.
|
chunk, unlike the vec0 table), so no dedup is needed here.
|
||||||
"""
|
"""
|
||||||
if not self.table_exists():
|
if not self.table_exists():
|
||||||
return {}
|
return {}
|
||||||
result: dict[str, str] = {}
|
return DocumentMetaTable.all_modified_times(self._conn)
|
||||||
for row in self._conn.execute(
|
|
||||||
"SELECT document_id, modified FROM " + DEFAULT_TABLE_NAME,
|
@property
|
||||||
):
|
def _db_path(self) -> str:
|
||||||
doc_id = str(row["document_id"])
|
return str(Path(self._uri) / DB_FILENAME)
|
||||||
if doc_id not in result:
|
|
||||||
result[doc_id] = str(row["modified"] or "")
|
@contextmanager
|
||||||
return result
|
def _rebuild_file(self) -> Iterator[sqlite3.Connection]:
|
||||||
|
"""Open a fresh temp database file for a file-swap rebuild (compact
|
||||||
|
or structural migration), yielding its connection for the caller to
|
||||||
|
populate.
|
||||||
|
|
||||||
|
On success, swaps the temp file in as the live database (closing
|
||||||
|
this store's current connection first -- see _swap_in_compact()).
|
||||||
|
On any exception, discards the temp file, including its -wal/-shm,
|
||||||
|
instead, and this store's own connection is left untouched.
|
||||||
|
"""
|
||||||
|
compact_path = self._db_path + ".compact"
|
||||||
|
new_conn = self._open_connection(compact_path)
|
||||||
|
try:
|
||||||
|
yield new_conn
|
||||||
|
except BaseException:
|
||||||
|
new_conn.close()
|
||||||
|
for suffix in ["", "-wal", "-shm"]:
|
||||||
|
Path(compact_path + suffix).unlink(missing_ok=True)
|
||||||
|
raise
|
||||||
|
else:
|
||||||
|
new_conn.close()
|
||||||
|
self._swap_in_compact(compact_path, self._db_path)
|
||||||
|
|
||||||
def compact(self, *, force: bool = False) -> None:
|
def compact(self, *, force: bool = False) -> None:
|
||||||
"""Rebuild the database file to reclaim space left behind by DELETEs.
|
"""Rebuild the database file to reclaim space left behind by DELETEs.
|
||||||
|
|
||||||
vec0 DELETE only invalidates rows; the vector data stays in the file
|
vec0 DELETE only invalidates rows; the vector data stays in the file
|
||||||
forever (asg017/sqlite-vec#54), and per-document re-indexing is a
|
forever, and per-document re-indexing is a delete+insert. The
|
||||||
delete+insert. The cumulative insert counter in ``index_meta`` tracks
|
cumulative insert counter in ``index_meta`` tracks total rows ever
|
||||||
total rows ever written; when that exceeds ``COMPACT_BLOAT_RATIO`` x
|
written; when that exceeds ``COMPACT_BLOAT_RATIO`` x the live row
|
||||||
the live row count (or when forced), live rows are copied into a fresh
|
count (or when forced), live rows are copied into a fresh database
|
||||||
database file and swapped in via ``os.replace``.
|
file and swapped in via ``os.replace``.
|
||||||
|
|
||||||
Note: ``ALTER TABLE ... RENAME TO`` on vec0 virtual tables does NOT
|
Note: ``ALTER TABLE ... RENAME TO`` on vec0 virtual tables does NOT
|
||||||
rename the shadow tables (sqlite-vec upstream limitation), so
|
rename the shadow tables (sqlite-vec upstream limitation), so an
|
||||||
an in-place rename-based rebuild is not safe. The file-swap approach
|
in-place rename-based rebuild is not safe. The file-swap approach is
|
||||||
is the maintainer-endorsed workaround (asg017/sqlite-vec#205).
|
the maintainer-endorsed workaround.
|
||||||
"""
|
"""
|
||||||
if not self.table_exists():
|
if not self.table_exists():
|
||||||
return
|
return
|
||||||
live = self._conn.execute(
|
if self.has_pending_migration():
|
||||||
"SELECT count(*) FROM " + DEFAULT_TABLE_NAME,
|
logger.warning(
|
||||||
).fetchone()[0]
|
"Skipping compact: store has a pending schema migration; "
|
||||||
total = int(self._meta_get("total_inserts") or str(live))
|
"run check_and_run_migrations() first",
|
||||||
|
)
|
||||||
|
return
|
||||||
|
live = DocumentChunksTable.count(self._conn)
|
||||||
|
total = IndexMetaTable.get_total_inserts(self._conn) or live
|
||||||
if not force and total <= max(live, 1) * COMPACT_BLOAT_RATIO:
|
if not force and total <= max(live, 1) * COMPACT_BLOAT_RATIO:
|
||||||
return
|
return
|
||||||
dim = self.vector_dim()
|
dim = self.vector_dim()
|
||||||
@@ -463,50 +574,61 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
live,
|
live,
|
||||||
total,
|
total,
|
||||||
)
|
)
|
||||||
db_path = str(Path(self._uri) / DB_FILENAME)
|
with self._rebuild_file() as new_conn:
|
||||||
compact_path = db_path + ".compact"
|
self._rebuild_into(self._conn, new_conn, dim)
|
||||||
|
|
||||||
# Copy all live rows into a fresh database file.
|
@staticmethod
|
||||||
new_conn = self._open_connection(compact_path)
|
def _rebuild_into(
|
||||||
try:
|
src_conn: sqlite3.Connection,
|
||||||
self._create_vec_table(new_conn, dim)
|
dst_conn: sqlite3.Connection,
|
||||||
self._meta_set_on(new_conn, "dim", str(dim))
|
dim: int,
|
||||||
for key in ("embed_model", "schema_version"):
|
) -> None:
|
||||||
value = self._meta_get(key)
|
"""Create the vec0 table in ``dst_conn``, copy dim/embed_model from
|
||||||
if value is not None:
|
``src_conn``, and stream every live vec0 row, document_chunks row,
|
||||||
self._meta_set_on(new_conn, key, value)
|
and document_meta row across. Used by compact() only --
|
||||||
src_cursor = self._conn.execute(
|
m0001_v1_to_v2 freezes its own copy loop instead of calling this,
|
||||||
"SELECT id, document_id, modified, node_content, embedding "
|
since this always reflects the *current* schema (see the migration
|
||||||
"FROM " + DEFAULT_TABLE_NAME,
|
DDL-freezing rule in the spec).
|
||||||
|
"""
|
||||||
|
PaperlessSqliteVecVectorStore._create_vec_table(dst_conn, dim)
|
||||||
|
dim_value = IndexMetaTable.get_dim(src_conn)
|
||||||
|
if dim_value is not None:
|
||||||
|
IndexMetaTable.set_dim(dst_conn, dim_value)
|
||||||
|
embed_model = IndexMetaTable.get_embed_model(src_conn)
|
||||||
|
if embed_model is not None:
|
||||||
|
IndexMetaTable.set_embed_model(dst_conn, embed_model)
|
||||||
|
schema_version = IndexMetaTable.get_schema_version(src_conn)
|
||||||
|
if schema_version is not None:
|
||||||
|
IndexMetaTable.set_schema_version(dst_conn, schema_version)
|
||||||
|
|
||||||
|
dst_conn.execute("BEGIN IMMEDIATE")
|
||||||
|
src_cursor = src_conn.execute(
|
||||||
|
"SELECT id, document_id, node_content, embedding FROM "
|
||||||
|
+ DEFAULT_TABLE_NAME,
|
||||||
|
)
|
||||||
|
copied = 0
|
||||||
|
while batch := src_cursor.fetchmany(BATCH_SIZE):
|
||||||
|
dst_conn.executemany(
|
||||||
|
_INSERT,
|
||||||
|
[
|
||||||
|
(
|
||||||
|
r["id"],
|
||||||
|
r["document_id"],
|
||||||
|
r["node_content"],
|
||||||
|
bytes(r["embedding"]),
|
||||||
|
)
|
||||||
|
for r in batch
|
||||||
|
],
|
||||||
)
|
)
|
||||||
new_conn.execute("BEGIN IMMEDIATE")
|
DocumentChunksTable.insert_many(
|
||||||
# Stream rows from the source cursor in batches instead of
|
dst_conn,
|
||||||
# materializing the whole table in memory, so a large index does
|
(ChunkRow(r["id"], r["document_id"]) for r in batch),
|
||||||
# not cause an OOM during routine maintenance compactions.
|
)
|
||||||
while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE):
|
copied += len(batch)
|
||||||
new_conn.executemany(
|
DocumentMetaTable.copy_all(src_conn, dst_conn, BATCH_SIZE)
|
||||||
self._INSERT,
|
# Reset the cumulative counter: after a rebuild, total_inserts == live.
|
||||||
[
|
IndexMetaTable.reset_total_inserts(dst_conn, copied)
|
||||||
(
|
dst_conn.execute("COMMIT")
|
||||||
r["id"],
|
|
||||||
r["document_id"],
|
|
||||||
r["modified"],
|
|
||||||
r["node_content"],
|
|
||||||
bytes(r["embedding"]),
|
|
||||||
)
|
|
||||||
for r in batch
|
|
||||||
],
|
|
||||||
)
|
|
||||||
# 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 _swap_in_compact(self, compact_path: str, db_path: str) -> None:
|
def _swap_in_compact(self, compact_path: str, db_path: str) -> None:
|
||||||
"""Atomically replace the live database with the compacted copy."""
|
"""Atomically replace the live database with the compacted copy."""
|
||||||
@@ -526,8 +648,8 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
"""
|
"""
|
||||||
if not self.table_exists():
|
if not self.table_exists():
|
||||||
return None
|
return None
|
||||||
raw = self._meta_get("schema_version")
|
raw_version = IndexMetaTable.get_schema_version(self._conn)
|
||||||
return int(raw) if raw is not None else SCHEMA_VERSION
|
return raw_version if raw_version is not None else SCHEMA_VERSION
|
||||||
|
|
||||||
def has_pending_migration(self) -> bool:
|
def has_pending_migration(self) -> bool:
|
||||||
"""Cheaply check whether a migration is pending, with no exclusive
|
"""Cheaply check whether a migration is pending, with no exclusive
|
||||||
@@ -591,16 +713,12 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
dim = self.vector_dim()
|
dim = self.vector_dim()
|
||||||
if dim is None: # pragma: no cover
|
if dim is None: # pragma: no cover
|
||||||
raise RuntimeError("Cannot migrate: no stored vector dimension")
|
raise RuntimeError("Cannot migrate: no stored vector dimension")
|
||||||
db_path = str(Path(self._uri) / DB_FILENAME)
|
with self._rebuild_file() as new_conn:
|
||||||
compact_path = db_path + ".compact"
|
|
||||||
new_conn = self._open_connection(compact_path)
|
|
||||||
try:
|
|
||||||
migration.apply(self._conn, new_conn, dim)
|
migration.apply(self._conn, new_conn, dim)
|
||||||
self._meta_set_on(new_conn, "schema_version", str(migration.to_version))
|
IndexMetaTable.set_schema_version(new_conn, migration.to_version)
|
||||||
except BaseException: # pragma: no cover
|
|
||||||
new_conn.close()
|
|
||||||
for p in [compact_path, compact_path + "-wal", compact_path + "-shm"]:
|
# Registers m0001_v1_to_v2 into MIGRATIONS; must be at the bottom (needs
|
||||||
Path(p).unlink(missing_ok=True)
|
# PaperlessSqliteVecVectorStore fully defined) -- see
|
||||||
raise
|
# paperless_ai/migrations/__init__.py for the full procedure.
|
||||||
new_conn.close()
|
from paperless_ai.migrations import m0001_v1_to_v2 # noqa: E402, F401
|
||||||
self._swap_in_compact(compact_path, db_path)
|
|
||||||
|
|||||||
Reference in New Issue
Block a user