mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-07-30 23:55:59 +00:00
Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
744e073ee6 | ||
|
|
40dbe3d903 | ||
|
|
f77ace3d13 | ||
|
|
f6cb5e3ef1 |
@@ -96,19 +96,26 @@ def get_context_for_document(
|
||||
user: User | None = None,
|
||||
max_docs: int = 5,
|
||||
) -> str:
|
||||
visible_documents = (
|
||||
get_objects_for_user_owner_aware(
|
||||
user,
|
||||
"view_document",
|
||||
Document,
|
||||
)
|
||||
if user
|
||||
else None
|
||||
)
|
||||
# None means "no restriction" to query_similar_documents. A superuser
|
||||
# (like no user at all) can see every document, so skip materializing
|
||||
# every visible pk into a Python list and passing it through as a SQL
|
||||
# IN filter: for a large library that is a wasted quadratic scan in the
|
||||
# vector store at best, and past ~32,763 documents a hard
|
||||
# sqlite3.OperationalError (SQLite's bound-parameter limit) at worst.
|
||||
# get_objects_for_user_owner_aware() would return every Document for a
|
||||
# superuser anyway (guardian's own with_superuser shortcut), so this
|
||||
# changes nothing about which documents are considered -- only how we
|
||||
# get there.
|
||||
visible_document_ids = (
|
||||
list(visible_documents.values_list("pk", flat=True))
|
||||
if visible_documents is not None
|
||||
else None
|
||||
None
|
||||
if user is None or user.is_superuser
|
||||
else list(
|
||||
get_objects_for_user_owner_aware(
|
||||
user,
|
||||
"view_document",
|
||||
Document,
|
||||
).values_list("pk", flat=True),
|
||||
)
|
||||
)
|
||||
similar_docs = query_similar_documents(
|
||||
document=doc,
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import enum
|
||||
import logging
|
||||
from collections.abc import Iterable
|
||||
from collections.abc import Iterator
|
||||
@@ -217,18 +218,49 @@ def write_store(embed_model_name: str | None = None):
|
||||
yield store
|
||||
|
||||
|
||||
def _check_and_run_migrations(store: "PaperlessSqliteVecVectorStore") -> bool:
|
||||
"""Run any pending structural migrations, returning True if a pending
|
||||
re-embed migration needs the caller to force a rebuild -- never
|
||||
triggered automatically here. Safe to call before any write, including
|
||||
class MigrationCheckResult(enum.Enum):
|
||||
"""Outcome of _check_and_run_migrations().
|
||||
|
||||
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)
|
||||
keeps this a no-op, with no exclusive access taken, once the store is
|
||||
current.
|
||||
"""
|
||||
if not store.has_pending_migration():
|
||||
return False
|
||||
return bool(
|
||||
_with_exclusive_access("migration check", store.check_and_run_migrations),
|
||||
return MigrationCheckResult.CURRENT
|
||||
result = _with_exclusive_access(
|
||||
"migration check",
|
||||
store.check_and_run_migrations,
|
||||
)
|
||||
if result is None:
|
||||
return MigrationCheckResult.DEFERRED
|
||||
return (
|
||||
MigrationCheckResult.REEMBED_REQUIRED
|
||||
if result
|
||||
else MigrationCheckResult.CURRENT
|
||||
)
|
||||
|
||||
|
||||
@@ -365,7 +397,7 @@ def _document_id_filters(doc_ids):
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.IN,
|
||||
value=sorted(doc_ids),
|
||||
value=list(doc_ids),
|
||||
),
|
||||
],
|
||||
)
|
||||
@@ -403,12 +435,21 @@ def update_llm_index(
|
||||
happens, since a rebuild always covers the whole library regardless.
|
||||
"""
|
||||
with write_store() as store:
|
||||
needs_reembed = _check_and_run_migrations(store)
|
||||
if needs_reembed:
|
||||
migration_result = _check_and_run_migrations(store)
|
||||
if migration_result is MigrationCheckResult.REEMBED_REQUIRED:
|
||||
logger.warning(
|
||||
"LLM index migration requires re-embedding; forcing rebuild.",
|
||||
)
|
||||
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(
|
||||
"correspondent",
|
||||
"document_type",
|
||||
@@ -482,8 +523,8 @@ def llm_index_add_or_update_document(document: Document):
|
||||
_embed_nodes(new_nodes, get_embedding_model(config))
|
||||
|
||||
with write_store(embed_model_name=get_configured_model_name(config)) as store:
|
||||
needs_reembed = _check_and_run_migrations(store)
|
||||
if needs_reembed:
|
||||
migration_result = _check_and_run_migrations(store)
|
||||
if migration_result is MigrationCheckResult.REEMBED_REQUIRED:
|
||||
logger.warning(
|
||||
"Skipping incremental LLM index update for document %s: the "
|
||||
"index requires re-embedding first. Run 'document_llmindex "
|
||||
@@ -491,6 +532,14 @@ def llm_index_add_or_update_document(document: Document):
|
||||
document.id,
|
||||
)
|
||||
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)
|
||||
|
||||
|
||||
@@ -509,14 +558,19 @@ def llm_index_migrate() -> None:
|
||||
if not AIConfig().llm_index_enabled:
|
||||
return
|
||||
with write_store() as store:
|
||||
needs_reembed = _check_and_run_migrations(store)
|
||||
if needs_reembed:
|
||||
migration_result = _check_and_run_migrations(store)
|
||||
if migration_result is MigrationCheckResult.REEMBED_REQUIRED:
|
||||
logger.warning(
|
||||
"LLM index requires re-embedding, which this automatic migration "
|
||||
"check will not do on its own -- it can be slow and, for a "
|
||||
"metered embedding backend, cost money. Run "
|
||||
"'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:
|
||||
@@ -528,7 +582,8 @@ def llm_index_compact() -> None:
|
||||
def llm_index_remove_document(document: Document):
|
||||
"""Remove a document's chunks from the LLM index."""
|
||||
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(
|
||||
"Skipping removal of document %s from the LLM index: the "
|
||||
"index requires re-embedding first. Run 'document_llmindex "
|
||||
@@ -536,6 +591,14 @@ def llm_index_remove_document(document: Document):
|
||||
document.id,
|
||||
)
|
||||
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))
|
||||
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@ 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 COMPACT_BATCH_SIZE
|
||||
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
|
||||
@@ -75,7 +75,7 @@ def _migrate_v1_to_v2(
|
||||
dst_conn.execute("BEGIN IMMEDIATE")
|
||||
src_cursor = src_conn.execute(_V1_SELECT)
|
||||
live = 0
|
||||
while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE):
|
||||
while batch := src_cursor.fetchmany(BATCH_SIZE):
|
||||
vec0_rows = []
|
||||
chunk_rows = []
|
||||
meta_by_document: dict[int, str] = {}
|
||||
|
||||
@@ -131,7 +131,7 @@ class DocumentMetaTable:
|
||||
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 COMPACT_BATCH_SIZE)."""
|
||||
it (pass BATCH_SIZE)."""
|
||||
cursor = src_conn.execute(
|
||||
"SELECT document_id, modified FROM document_meta",
|
||||
)
|
||||
|
||||
@@ -3,6 +3,8 @@ from unittest.mock import MagicMock
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from django.contrib.auth.models import User
|
||||
from django.test import override_settings
|
||||
|
||||
from documents.models import Document
|
||||
@@ -278,3 +280,104 @@ def test_get_context_for_document_no_similar_docs(mock_document):
|
||||
with patch("paperless_ai.ai_classifier.query_similar_documents", return_value=[]):
|
||||
result = get_context_for_document(mock_document)
|
||||
assert result == ""
|
||||
|
||||
|
||||
class TestGetContextForDocumentVisibility:
|
||||
"""get_context_for_document must not materialize every visible document
|
||||
id for a user who can already see the whole library: a superuser (like
|
||||
no user at all) gets document_ids=None (no restriction) straight
|
||||
through to query_similar_documents(), instead of a full-library IN
|
||||
filter that is wasteful at best and, past ~32,763 documents, a hard
|
||||
sqlite3.OperationalError at worst (SQLite's bound-parameter limit).
|
||||
"""
|
||||
|
||||
def test_skips_permission_lookup_for_superuser(
|
||||
self,
|
||||
mock_document: MagicMock,
|
||||
mock_similar_documents: list[MagicMock],
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- A superuser
|
||||
WHEN:
|
||||
- get_context_for_document() is called
|
||||
THEN:
|
||||
- get_objects_for_user_owner_aware() is never called, and
|
||||
query_similar_documents() is called with document_ids=None
|
||||
"""
|
||||
mock_query = mocker.patch(
|
||||
"paperless_ai.ai_classifier.query_similar_documents",
|
||||
return_value=mock_similar_documents,
|
||||
)
|
||||
mock_get_objects = mocker.patch(
|
||||
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
|
||||
)
|
||||
user = mocker.MagicMock(spec=User)
|
||||
user.is_superuser = True
|
||||
|
||||
get_context_for_document(mock_document, user, max_docs=2)
|
||||
|
||||
mock_get_objects.assert_not_called()
|
||||
assert mock_query.call_args.kwargs["document_ids"] is None
|
||||
|
||||
def test_skips_permission_lookup_when_no_user(
|
||||
self,
|
||||
mock_document: MagicMock,
|
||||
mock_similar_documents: list[MagicMock],
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- No user (user=None)
|
||||
WHEN:
|
||||
- get_context_for_document() is called
|
||||
THEN:
|
||||
- get_objects_for_user_owner_aware() is never called, and
|
||||
query_similar_documents() is called with document_ids=None
|
||||
"""
|
||||
mock_query = mocker.patch(
|
||||
"paperless_ai.ai_classifier.query_similar_documents",
|
||||
return_value=mock_similar_documents,
|
||||
)
|
||||
mock_get_objects = mocker.patch(
|
||||
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
|
||||
)
|
||||
|
||||
get_context_for_document(mock_document, None, max_docs=2)
|
||||
|
||||
mock_get_objects.assert_not_called()
|
||||
assert mock_query.call_args.kwargs["document_ids"] is None
|
||||
|
||||
def test_restricts_to_visible_documents_for_non_superuser(
|
||||
self,
|
||||
mock_document: MagicMock,
|
||||
mock_similar_documents: list[MagicMock],
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- A non-superuser with a specific set of visible documents
|
||||
WHEN:
|
||||
- get_context_for_document() is called
|
||||
THEN:
|
||||
- query_similar_documents() is called with exactly that user's
|
||||
visible document ids, unchanged from before this optimization
|
||||
"""
|
||||
mock_query = mocker.patch(
|
||||
"paperless_ai.ai_classifier.query_similar_documents",
|
||||
return_value=mock_similar_documents,
|
||||
)
|
||||
mock_queryset = mocker.MagicMock()
|
||||
mock_queryset.values_list.return_value = [1, 2, 3]
|
||||
mock_get_objects = mocker.patch(
|
||||
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
|
||||
return_value=mock_queryset,
|
||||
)
|
||||
user = mocker.MagicMock(spec=User)
|
||||
user.is_superuser = False
|
||||
|
||||
get_context_for_document(mock_document, user, max_docs=2)
|
||||
|
||||
mock_get_objects.assert_called_once_with(user, "view_document", Document)
|
||||
assert mock_query.call_args.kwargs["document_ids"] == [1, 2, 3]
|
||||
|
||||
@@ -9,6 +9,7 @@ from django.db import connection
|
||||
from django.test import override_settings
|
||||
from django.test.utils import CaptureQueriesContext
|
||||
from django.utils import timezone
|
||||
from filelock import Timeout
|
||||
from llama_index.core.schema import MetadataMode
|
||||
|
||||
from documents.models import Correspondent
|
||||
@@ -869,6 +870,44 @@ class TestLlmIndexLocking:
|
||||
|
||||
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(
|
||||
self,
|
||||
temp_llm_index_dir: Path,
|
||||
@@ -915,6 +954,34 @@ class TestLlmIndexLocking:
|
||||
|
||||
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(
|
||||
self,
|
||||
temp_llm_index_dir: Path,
|
||||
@@ -938,6 +1005,35 @@ class TestLlmIndexLocking:
|
||||
|
||||
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
|
||||
@@ -1056,6 +1152,37 @@ class TestLlmIndexMigrate:
|
||||
indexing.llm_index_migrate()
|
||||
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
|
||||
class TestQuerySimilarDocuments:
|
||||
|
||||
@@ -17,6 +17,7 @@ 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 _MAX_IN_VALUES
|
||||
from paperless_ai.vector_store import DB_FILENAME
|
||||
from paperless_ai.vector_store import DEFAULT_TABLE_NAME
|
||||
from paperless_ai.vector_store import SCHEMA_VERSION
|
||||
@@ -296,8 +297,47 @@ class TestBuildWhere:
|
||||
assert where == "1 = 0"
|
||||
assert params == []
|
||||
|
||||
def test_query_with_untranslatable_filter_returns_no_rows(self, store) -> None:
|
||||
store.add([make_node("a1", 1), make_node("b1", 2)])
|
||||
def test_fails_closed_when_in_filter_exceeds_max_values(
|
||||
self,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- An IN filter with more values than _MAX_IN_VALUES (SQLite's
|
||||
own bound-parameter limit is 32766; this guard sits below
|
||||
that with headroom for the query's other bound parameters)
|
||||
WHEN:
|
||||
- _build_where() translates it to SQL
|
||||
THEN:
|
||||
- It fails closed ("1 = 0", no params) instead of building an
|
||||
IN clause SQLite would reject, and logs a warning -- this
|
||||
filter scopes document access, so refusing to build it must
|
||||
never widen the scope to "everything" by accident
|
||||
"""
|
||||
oversized = _in_filter([str(i) for i in range(_MAX_IN_VALUES + 1)])
|
||||
|
||||
with caplog.at_level("WARNING"):
|
||||
where, params = _build_where(oversized)
|
||||
|
||||
assert where == "(1 = 0)"
|
||||
assert params == []
|
||||
assert "document_id" in caplog.text
|
||||
|
||||
def test_query_with_untranslatable_filter_returns_no_rows(
|
||||
self,
|
||||
store: PaperlessSqliteVecVectorStore,
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- A store with one chunk each for documents 1 and 2
|
||||
WHEN:
|
||||
- A query and a get_nodes() call are issued with a filter whose
|
||||
only clause is an untranslatable nested MetadataFilters
|
||||
THEN:
|
||||
- Neither call raises, and both fail closed by returning no rows
|
||||
(rather than emitting invalid or unfiltered SQL)
|
||||
"""
|
||||
store.add([make_node("a1", "1"), make_node("b1", "2")])
|
||||
nested = MetadataFilters(
|
||||
filters=[
|
||||
MetadataFilter(
|
||||
@@ -508,7 +548,7 @@ class TestCompact:
|
||||
A tiny batch size forces several fetchmany()/executemany() cycles so a
|
||||
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.compact(force=True)
|
||||
ids = {n.node_id for n in store.get_nodes(filters=_in_filter([1]))}
|
||||
|
||||
@@ -47,10 +47,12 @@ SCHEMA_VERSION = 2
|
||||
# a rebuild copies the live rows into a fresh table.
|
||||
COMPACT_BLOAT_RATIO = 2.0
|
||||
|
||||
# compact(): number of rows copied per executemany() when rebuilding the file.
|
||||
# Rows are streamed from the source cursor in batches of this size rather than
|
||||
# materialized all at once, keeping memory bounded regardless of index size.
|
||||
COMPACT_BATCH_SIZE = 500
|
||||
# Number of rows fetched/copied per batch whenever this module streams rows
|
||||
# instead of materializing them all at once, keeping memory bounded regardless
|
||||
# of index size -- used by compact()'s rebuild, m0001_v1_to_v2's migration
|
||||
# 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
|
||||
# keys we construct ourselves, but allowlisting keeps SQL identifiers safe by
|
||||
@@ -73,6 +75,17 @@ class _Row(NamedTuple):
|
||||
embedding: bytes
|
||||
|
||||
|
||||
# _build_where(): the largest IN value list translated into bound SQL
|
||||
# parameters. SQLite's own hard limit (SQLITE_MAX_VARIABLE_NUMBER) is 32766
|
||||
# by default; this leaves headroom below that for the query's other bound
|
||||
# parameters (the embedding blob, k, and any NE clause) and for the limit
|
||||
# itself to move. An IN filter this large should not happen in practice --
|
||||
# callers are expected to pass None (no filter) rather than every id when
|
||||
# the filter would not actually narrow anything -- so this is a guard
|
||||
# against a future regression, not a normal code path.
|
||||
_MAX_IN_VALUES = 32700
|
||||
|
||||
|
||||
def _pack(embedding: Sequence[float]) -> bytes:
|
||||
return struct.pack(f"{len(embedding)}f", *embedding)
|
||||
|
||||
@@ -117,6 +130,21 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]:
|
||||
if not values: # pragma: no cover
|
||||
clauses.append("1 = 0")
|
||||
continue
|
||||
if len(values) > _MAX_IN_VALUES:
|
||||
# Fail closed (see the empty-clauses case below) rather than
|
||||
# let SQLite raise "too many SQL variables" past its own
|
||||
# limit: this filter scopes document access, so an IN list
|
||||
# too large to safely bind must match no rows, never widen
|
||||
# the scope to "everything" by accident.
|
||||
logger.warning(
|
||||
"Refusing to build an IN filter on %r with %d values "
|
||||
"(over the %d-value safety limit); returning no rows.",
|
||||
f.key,
|
||||
len(values),
|
||||
_MAX_IN_VALUES,
|
||||
)
|
||||
clauses.append("1 = 0")
|
||||
continue
|
||||
placeholders = ",".join("?" for _ in values)
|
||||
clauses.append(f"{f.key} IN ({placeholders})")
|
||||
params.extend(values)
|
||||
@@ -605,7 +633,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
||||
+ DEFAULT_TABLE_NAME,
|
||||
)
|
||||
copied = 0
|
||||
while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE):
|
||||
while batch := src_cursor.fetchmany(BATCH_SIZE):
|
||||
dst_conn.executemany(
|
||||
_INSERT,
|
||||
[
|
||||
@@ -623,7 +651,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
||||
(ChunkRow(r["id"], r["document_id"]) for r in batch),
|
||||
)
|
||||
copied += len(batch)
|
||||
DocumentMetaTable.copy_all(src_conn, dst_conn, COMPACT_BATCH_SIZE)
|
||||
DocumentMetaTable.copy_all(src_conn, dst_conn, BATCH_SIZE)
|
||||
# Reset the cumulative counter: after a rebuild, total_inserts == live.
|
||||
IndexMetaTable.reset_total_inserts(dst_conn, copied)
|
||||
dst_conn.execute("COMMIT")
|
||||
|
||||
Reference in New Issue
Block a user