mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-09-03 16:35:06 +00:00
Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2f440cade2 | ||
|
|
61397749b7 | ||
|
|
3666234f51 | ||
|
|
ec7745e71f | ||
|
|
98e87d91ad |
+1
-1
@@ -247,7 +247,7 @@ per-file-ignores."src/documents/models.py" = [
|
||||
isort.force-single-line = true
|
||||
|
||||
[tool.codespell]
|
||||
ignore-words-list = "criterias,afterall,valeu,ureue,equest,ure,assertIn,Oktober,commitish"
|
||||
ignore-words-list = "criterias,afterall,valeu,ureue,equest,ure,assertIn,Oktober,commitish,NIN,nin"
|
||||
skip = """\
|
||||
src-ui/src/locale/*,src-ui/pnpm-lock.yaml,src-ui/e2e/*,src/paperless_mail/tests/samples/*,src/paperless/tests/samples\
|
||||
/mail/*,src/documents/tests/samples/*,*.po,*.json\
|
||||
|
||||
@@ -1,11 +1,18 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from django.contrib.auth.models import Permission
|
||||
from django.contrib.auth.models import User
|
||||
from rest_framework import status
|
||||
from rest_framework.test import APIClient
|
||||
from rest_framework.test import APITestCase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
|
||||
class TestChatStreamingViewInputValidation(APITestCase):
|
||||
def setUp(self) -> None:
|
||||
@@ -42,3 +49,73 @@ class TestChatStreamingViewInputValidation(APITestCase):
|
||||
format="json",
|
||||
)
|
||||
assert resp.status_code == status.HTTP_400_BAD_REQUEST
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
class TestChatStreamingViewUnrestrictedFlag:
|
||||
"""The document id filter may only be skipped (``unrestricted=True``) for
|
||||
an active superuser, never for a regular user -- regardless of what
|
||||
permissions that user holds.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def mocked_stream_chat(self, mocker: MockerFixture) -> mock.MagicMock:
|
||||
"""AI enabled, with stream_chat_with_documents patched so the view
|
||||
never touches the real vector store; returns the patched callable so
|
||||
tests can inspect how it was called.
|
||||
"""
|
||||
mocker.patch("documents.views.AIConfig").return_value.ai_enabled = True
|
||||
return mocker.patch(
|
||||
"documents.views.stream_chat_with_documents",
|
||||
return_value=iter(()),
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def viewer_client(self, user_client: APIClient, regular_user: User) -> APIClient:
|
||||
"""The conftest regular-user client, granted the global
|
||||
view_document permission -- the minimum ViewDocumentsPermissions
|
||||
needs to reach the view at all. Model-level only: says nothing
|
||||
about which documents (if any) this user can actually see.
|
||||
"""
|
||||
regular_user.user_permissions.add(
|
||||
*Permission.objects.filter(codename="view_document"),
|
||||
)
|
||||
return user_client
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("client_fixture", "expected_unrestricted"),
|
||||
[
|
||||
pytest.param("admin_client", True, id="superuser_is_unrestricted"),
|
||||
pytest.param("viewer_client", False, id="regular_user_is_restricted"),
|
||||
],
|
||||
)
|
||||
def test_unrestricted_only_for_superuser(
|
||||
self,
|
||||
request: pytest.FixtureRequest,
|
||||
mocked_stream_chat: mock.MagicMock,
|
||||
client_fixture: str,
|
||||
*,
|
||||
expected_unrestricted: bool,
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- A superuser, or a regular user holding the global
|
||||
view_document permission (but no object-level document access)
|
||||
WHEN:
|
||||
- They post a chat question with no document_id
|
||||
THEN:
|
||||
- stream_chat_with_documents is called with unrestricted=True
|
||||
only for the superuser; the regular user is always
|
||||
unrestricted=False, regardless of their permissions
|
||||
"""
|
||||
client: APIClient = request.getfixturevalue(client_fixture)
|
||||
|
||||
client.post(
|
||||
"/api/documents/chat/",
|
||||
data={"q": "What's in these documents?"},
|
||||
format="json",
|
||||
)
|
||||
|
||||
assert (
|
||||
mocked_stream_chat.call_args.kwargs["unrestricted"] is expected_unrestricted
|
||||
)
|
||||
|
||||
@@ -180,6 +180,7 @@ from documents.permissions import has_system_status_permission
|
||||
from documents.permissions import permitted_document_ids
|
||||
from documents.permissions import permitted_object_ids
|
||||
from documents.permissions import set_permissions_for_object
|
||||
from documents.permissions import user_is_unrestricted
|
||||
from documents.plugins.date_parsing import get_date_parser
|
||||
from documents.schema import generate_object_with_permissions_schema
|
||||
from documents.search import SearchHit
|
||||
@@ -2329,10 +2330,12 @@ class ChatStreamingView(GenericAPIView[Any]):
|
||||
return HttpResponseForbidden("Insufficient permissions")
|
||||
|
||||
documents = Document.objects.filter(pk=document.pk)
|
||||
unrestricted = False
|
||||
else:
|
||||
documents = Document.objects.filter(
|
||||
id__in=permitted_document_ids(request.user),
|
||||
)
|
||||
unrestricted = user_is_unrestricted(request.user)
|
||||
|
||||
output_language = get_llm_output_language(
|
||||
ai_config=ai_config,
|
||||
@@ -2343,6 +2346,7 @@ class ChatStreamingView(GenericAPIView[Any]):
|
||||
stream_chat_with_documents(
|
||||
query_str=question,
|
||||
documents=documents,
|
||||
unrestricted=unrestricted,
|
||||
output_language=output_language,
|
||||
),
|
||||
content_type="text/event-stream",
|
||||
|
||||
@@ -8,7 +8,8 @@ from documents.models import Document
|
||||
from paperless.config import AIConfig
|
||||
from paperless_ai.client import AIClient
|
||||
from paperless_ai.db import db_connection_released
|
||||
from paperless_ai.indexing import _document_id_filters
|
||||
from paperless_ai.indexing import document_id_filters
|
||||
from paperless_ai.indexing import exclude_document_ids_filter
|
||||
from paperless_ai.indexing import get_rag_prompt_helper
|
||||
from paperless_ai.indexing import load_or_build_index
|
||||
from paperless_ai.indexing import read_store
|
||||
@@ -95,12 +96,15 @@ def _format_chat_metadata_trailer(references: list[dict[str, int | str]]) -> str
|
||||
def stream_chat_with_documents(
|
||||
query_str: str,
|
||||
documents: QuerySet[Document],
|
||||
*,
|
||||
unrestricted: bool = False,
|
||||
output_language: str | None = None,
|
||||
):
|
||||
try:
|
||||
yield from _stream_chat_with_documents(
|
||||
query_str,
|
||||
documents,
|
||||
unrestricted=unrestricted,
|
||||
output_language=output_language,
|
||||
)
|
||||
except Exception as e:
|
||||
@@ -111,6 +115,8 @@ def stream_chat_with_documents(
|
||||
def _stream_chat_with_documents(
|
||||
query_str: str,
|
||||
documents: QuerySet[Document],
|
||||
*,
|
||||
unrestricted: bool = False,
|
||||
output_language: str | None = None,
|
||||
):
|
||||
if not documents.exists():
|
||||
@@ -123,9 +129,20 @@ def _stream_chat_with_documents(
|
||||
from llama_index.core.retrievers import VectorIndexRetriever
|
||||
|
||||
config = AIConfig()
|
||||
filters = _document_id_filters(
|
||||
str(pk) for pk in documents.values_list("pk", flat=True)
|
||||
)
|
||||
if unrestricted:
|
||||
# Exclude trashed ids (usually few) instead of an IN filter over the
|
||||
# full permitted set, which risks the vector store's bound parameter
|
||||
# limit (_MAX_IN_VALUES) on large installs. Trashed documents stay
|
||||
# indexed until permanent deletion (delete_document_from_llm_index
|
||||
# hangs off post_delete, not trash), so must be excluded explicitly.
|
||||
trashed_ids = Document.global_objects.filter(
|
||||
deleted_at__isnull=False,
|
||||
).values_list("pk", flat=True)
|
||||
filters = exclude_document_ids_filter(str(pk) for pk in trashed_ids)
|
||||
else:
|
||||
filters = document_id_filters(
|
||||
str(pk) for pk in documents.values_list("pk", flat=True)
|
||||
)
|
||||
|
||||
# Hold the shared read lock for the whole operation: the query engine
|
||||
# retrieves from the vector store again during synthesis, so the connection
|
||||
|
||||
@@ -362,7 +362,7 @@ def _embed_nodes(nodes: list["BaseNode"], embed_model) -> None:
|
||||
node.embedding = emb
|
||||
|
||||
|
||||
def _document_id_filters(doc_ids):
|
||||
def document_id_filters(doc_ids):
|
||||
"""Return a MetadataFilters IN filter scoped to ``doc_ids``."""
|
||||
from llama_index.core.vector_stores.types import FilterOperator
|
||||
from llama_index.core.vector_stores.types import MetadataFilter
|
||||
@@ -396,6 +396,23 @@ def _exclude_document_id_filter(document_id: int | str):
|
||||
)
|
||||
|
||||
|
||||
def exclude_document_ids_filter(doc_ids):
|
||||
"""Return a MetadataFilters NIN filter excluding every id in ``doc_ids``."""
|
||||
from llama_index.core.vector_stores.types import FilterOperator
|
||||
from llama_index.core.vector_stores.types import MetadataFilter
|
||||
from llama_index.core.vector_stores.types import MetadataFilters
|
||||
|
||||
return MetadataFilters(
|
||||
filters=[
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.NIN,
|
||||
value=list(doc_ids),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def update_llm_index(
|
||||
*,
|
||||
iter_wrapper: IterWrapper[Document] = identity,
|
||||
@@ -660,7 +677,7 @@ def retrieve_similar_nodes(
|
||||
|
||||
filter_parts = []
|
||||
if allowed_document_ids is not None:
|
||||
filter_parts.extend(_document_id_filters(allowed_document_ids).filters)
|
||||
filter_parts.extend(document_id_filters(allowed_document_ids).filters)
|
||||
if document.pk is not None:
|
||||
filter_parts.extend(_exclude_document_id_filter(document.pk).filters)
|
||||
|
||||
|
||||
@@ -154,6 +154,35 @@ class DocumentMetaTable:
|
||||
}
|
||||
|
||||
|
||||
class PermittedIdsTable:
|
||||
"""Per-connection scratch space for an oversized IN-filter id list.
|
||||
|
||||
A literal ``IN (?,?,...)`` list binds one SQL parameter per id, capped by
|
||||
SQLite's own SQLITE_MAX_VARIABLE_NUMBER (see _MAX_IN_VALUES in
|
||||
vector_store.py). Loading the ids into a TEMP TABLE and filtering via a
|
||||
subquery instead has no such limit. TEMP tables live in a
|
||||
connection-private namespace -- never visible to another connection,
|
||||
even under this identical name -- so this is safe under the vector
|
||||
store's one-connection-per-request model without any extra locking or
|
||||
per-call naming scheme.
|
||||
"""
|
||||
|
||||
TABLE_NAME = "permitted_document_ids"
|
||||
|
||||
@staticmethod
|
||||
def load(conn: sqlite3.Connection, ids: Iterable[int]) -> None:
|
||||
"""Replace this connection's scratch table with ``ids``."""
|
||||
conn.execute(f"DROP TABLE IF EXISTS temp.{PermittedIdsTable.TABLE_NAME}")
|
||||
conn.execute(
|
||||
f"CREATE TEMP TABLE {PermittedIdsTable.TABLE_NAME} "
|
||||
"(id INTEGER PRIMARY KEY)",
|
||||
)
|
||||
conn.executemany(
|
||||
f"INSERT INTO {PermittedIdsTable.TABLE_NAME} (id) VALUES (?)",
|
||||
((i,) for i in ids),
|
||||
)
|
||||
|
||||
|
||||
class IndexMetaTable:
|
||||
"""Typed accessors over index_meta's key/value rows -- replaces
|
||||
PaperlessSqliteVecVectorStore._meta_get_on/_meta_set_on, which returned
|
||||
|
||||
@@ -1,9 +1,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from django.db.models.signals import post_init
|
||||
from django.utils import timezone
|
||||
from llama_index.core import settings as llama_settings
|
||||
from llama_index.core.embeddings.mock_embed_model import MockEmbedding
|
||||
from llama_index.core.schema import TextNode
|
||||
@@ -18,6 +23,11 @@ from paperless_ai.chat import _build_chat_prompt
|
||||
from paperless_ai.chat import _build_refine_prompt
|
||||
from paperless_ai.chat import stream_chat_with_documents
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
import pytest_mock
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_embed_model():
|
||||
@@ -310,8 +320,40 @@ def test_stream_chat_unexpected_failure_returns_generic_error(caplog) -> None:
|
||||
assert "private provider detail" in caplog.text
|
||||
|
||||
|
||||
def _retriever_filter_values(captured_filters: list[Any]) -> list[str]:
|
||||
"""The value list of the single MetadataFilter the retriever received."""
|
||||
assert captured_filters, "VectorIndexRetriever was never constructed"
|
||||
filt = captured_filters[0]
|
||||
assert filt is not None, "Retriever must receive a MetadataFilters"
|
||||
return filt.filters[0].value
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
class TestStreamChatRetrieval:
|
||||
@pytest.fixture
|
||||
def captured_filters(self, mocker: pytest_mock.MockerFixture) -> list[Any]:
|
||||
"""Stub out the AI client and the retriever, capturing the ``filters``
|
||||
kwarg of every VectorIndexRetriever construction.
|
||||
|
||||
VectorIndexRetriever is imported inside _stream_chat_with_documents,
|
||||
so it is patched at the llama_index source for the lazy import to
|
||||
pick it up.
|
||||
"""
|
||||
captured: list[Any] = []
|
||||
retriever = mocker.MagicMock()
|
||||
retriever.retrieve.return_value = []
|
||||
|
||||
def capture_retriever(*args, **kwargs) -> pytest_mock.MockType:
|
||||
captured.append(kwargs.get("filters"))
|
||||
return retriever
|
||||
|
||||
mocker.patch("paperless_ai.chat.AIClient")
|
||||
mocker.patch(
|
||||
"llama_index.core.retrievers.VectorIndexRetriever",
|
||||
side_effect=capture_retriever,
|
||||
)
|
||||
return captured
|
||||
|
||||
def test_no_nodes_yields_no_content_message(
|
||||
self,
|
||||
temp_llm_index_dir,
|
||||
@@ -329,9 +371,9 @@ class TestStreamChatRetrieval:
|
||||
|
||||
def test_chat_filter_contains_only_requested_document_ids(
|
||||
self,
|
||||
temp_llm_index_dir,
|
||||
mock_embed_model,
|
||||
mocker,
|
||||
temp_llm_index_dir: Path,
|
||||
mock_embed_model: pytest_mock.MockType,
|
||||
captured_filters: list[Any],
|
||||
) -> None:
|
||||
"""The MetadataFilter passed to the retriever must be scoped to the
|
||||
requested documents only — content from other indexed documents must
|
||||
@@ -342,22 +384,6 @@ class TestStreamChatRetrieval:
|
||||
indexing.llm_index_add_or_update_document(included)
|
||||
indexing.llm_index_add_or_update_document(excluded)
|
||||
|
||||
# VectorIndexRetriever is imported inside _stream_chat_with_documents;
|
||||
# patch it at the llama_index source so the lazy import picks it up.
|
||||
captured_filters = []
|
||||
mock_retriever = mocker.MagicMock()
|
||||
mock_retriever.retrieve.return_value = []
|
||||
|
||||
def capture_retriever(*args, **kwargs):
|
||||
captured_filters.append(kwargs.get("filters"))
|
||||
return mock_retriever
|
||||
|
||||
mocker.patch("paperless_ai.chat.AIClient")
|
||||
mocker.patch(
|
||||
"llama_index.core.retrievers.VectorIndexRetriever",
|
||||
side_effect=capture_retriever,
|
||||
)
|
||||
|
||||
list(
|
||||
chat.stream_chat_with_documents(
|
||||
"question?",
|
||||
@@ -365,13 +391,78 @@ class TestStreamChatRetrieval:
|
||||
),
|
||||
)
|
||||
|
||||
assert captured_filters, "VectorIndexRetriever was never constructed"
|
||||
filt = captured_filters[0]
|
||||
assert filt is not None, "Retriever must receive a MetadataFilters"
|
||||
filter_values = filt.filters[0].value
|
||||
filter_values = _retriever_filter_values(captured_filters)
|
||||
assert str(included.pk) in filter_values
|
||||
assert str(excluded.pk) not in filter_values
|
||||
|
||||
def test_unrestricted_chat_excludes_nothing_when_no_documents_are_trashed(
|
||||
self,
|
||||
temp_llm_index_dir: Path,
|
||||
mock_embed_model: pytest_mock.MockType,
|
||||
captured_filters: list[Any],
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- A document indexed in the vector store, nothing trashed
|
||||
WHEN:
|
||||
- stream_chat_with_documents is called with unrestricted=True
|
||||
THEN:
|
||||
- The retriever receives a NOT IN filter excluding zero ids, so
|
||||
the whole index is effectively searched -- and no IN-list is
|
||||
built from the full permitted set, which is what risks the
|
||||
vector store's safety limit on large installs
|
||||
"""
|
||||
document = DocumentFactory.create(content="indexed document content")
|
||||
indexing.llm_index_add_or_update_document(document)
|
||||
|
||||
list(
|
||||
chat.stream_chat_with_documents(
|
||||
"question?",
|
||||
Document.objects.filter(pk=document.pk),
|
||||
unrestricted=True,
|
||||
),
|
||||
)
|
||||
|
||||
assert _retriever_filter_values(captured_filters) == []
|
||||
|
||||
def test_unrestricted_chat_excludes_trashed_documents(
|
||||
self,
|
||||
temp_llm_index_dir: Path,
|
||||
mock_embed_model: pytest_mock.MockType,
|
||||
captured_filters: list[Any],
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- Two indexed documents, one of them trashed -- trashed documents
|
||||
stay in the vector index until permanently deleted, since
|
||||
delete_document_from_llm_index is wired to post_delete
|
||||
WHEN:
|
||||
- stream_chat_with_documents is called with unrestricted=True
|
||||
THEN:
|
||||
- The retriever receives a NOT IN filter excluding the trashed
|
||||
document's id, so an unrestricted caller (e.g. a superuser)
|
||||
never has trashed content surfaced in a chat answer
|
||||
"""
|
||||
kept = DocumentFactory.create(content="kept document content")
|
||||
trashed = DocumentFactory.create(content="trashed document content")
|
||||
indexing.llm_index_add_or_update_document(kept)
|
||||
indexing.llm_index_add_or_update_document(trashed)
|
||||
Document.global_objects.filter(pk=trashed.pk).update(
|
||||
deleted_at=timezone.now(),
|
||||
)
|
||||
|
||||
list(
|
||||
chat.stream_chat_with_documents(
|
||||
"question?",
|
||||
Document.objects.filter(pk=kept.pk),
|
||||
unrestricted=True,
|
||||
),
|
||||
)
|
||||
|
||||
filter_values = _retriever_filter_values(captured_filters)
|
||||
assert str(trashed.pk) in filter_values
|
||||
assert str(kept.pk) not in filter_values
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_get_document_references_only_queries_referenced_documents(
|
||||
self,
|
||||
|
||||
@@ -9,6 +9,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.tables import PermittedIdsTable
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -338,3 +339,85 @@ class TestIndexMetaTable:
|
||||
IndexMetaTable.increment_total_inserts(conn, 100)
|
||||
IndexMetaTable.reset_total_inserts(conn, 7)
|
||||
assert IndexMetaTable.get_total_inserts(conn) == 7
|
||||
|
||||
|
||||
class TestPermittedIdsTable:
|
||||
def _loaded_ids(self, conn: sqlite3.Connection) -> list[int]:
|
||||
return [
|
||||
row["id"]
|
||||
for row in conn.execute(
|
||||
f"SELECT id FROM {PermittedIdsTable.TABLE_NAME} ORDER BY id",
|
||||
)
|
||||
]
|
||||
|
||||
def test_load_then_read_back_all_ids(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- A bare sqlite3 connection
|
||||
WHEN:
|
||||
- load() is called with a set of ids
|
||||
THEN:
|
||||
- Every id is present in the TEMP TABLE, and only those ids
|
||||
"""
|
||||
PermittedIdsTable.load(conn, [3, 1, 2])
|
||||
assert self._loaded_ids(conn) == [1, 2, 3]
|
||||
|
||||
def test_load_replaces_previous_contents(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- A connection whose PermittedIdsTable already holds one id set
|
||||
WHEN:
|
||||
- load() is called again with a different id set
|
||||
THEN:
|
||||
- Only the new ids are present -- a connection reused across
|
||||
multiple queries in one request never leaks a stale filter
|
||||
"""
|
||||
PermittedIdsTable.load(conn, [1, 2, 3])
|
||||
PermittedIdsTable.load(conn, [4, 5])
|
||||
assert self._loaded_ids(conn) == [4, 5]
|
||||
|
||||
def test_load_is_connection_private(self) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- Two separate connections
|
||||
WHEN:
|
||||
- Each loads PermittedIdsTable with a different id set, under
|
||||
the identical TABLE_NAME
|
||||
THEN:
|
||||
- Each connection sees only its own ids -- TEMP TABLE is
|
||||
connection-private, so concurrent requests never collide or
|
||||
cross-contaminate despite sharing the same table name (the
|
||||
vector store opens one connection per request; see
|
||||
PaperlessSqliteVecVectorStore)
|
||||
"""
|
||||
conn_a = sqlite3.connect(":memory:")
|
||||
conn_a.row_factory = sqlite3.Row
|
||||
conn_b = sqlite3.connect(":memory:")
|
||||
conn_b.row_factory = sqlite3.Row
|
||||
try:
|
||||
PermittedIdsTable.load(conn_a, [1, 2, 3])
|
||||
PermittedIdsTable.load(conn_b, [4, 5, 6])
|
||||
assert self._loaded_ids(conn_a) == [1, 2, 3]
|
||||
assert self._loaded_ids(conn_b) == [4, 5, 6]
|
||||
finally:
|
||||
conn_a.close()
|
||||
conn_b.close()
|
||||
|
||||
def test_load_handles_more_ids_than_a_bound_parameter_list_could(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- An id count over SQLite's own bound-parameter limit
|
||||
(SQLITE_MAX_VARIABLE_NUMBER, 32766 by default) -- more than a
|
||||
literal IN(?,?,...) list could ever bind in one statement
|
||||
WHEN:
|
||||
- load() is called with that many ids
|
||||
THEN:
|
||||
- Every id is loaded without error, since executemany() binds
|
||||
one row at a time rather than one statement with N parameters
|
||||
"""
|
||||
ids = list(range(40_000))
|
||||
PermittedIdsTable.load(conn, ids)
|
||||
assert self._loaded_ids(conn) == ids
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import inspect
|
||||
import sqlite3
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Generator
|
||||
from pathlib import Path
|
||||
|
||||
@@ -17,6 +18,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.tables import PermittedIdsTable
|
||||
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
|
||||
@@ -97,6 +99,18 @@ def _ne_filter(document_id: int):
|
||||
)
|
||||
|
||||
|
||||
def _nin_filter(document_ids: list[int]):
|
||||
return MetadataFilters(
|
||||
filters=[
|
||||
MetadataFilter(
|
||||
key="document_id",
|
||||
operator=FilterOperator.NIN,
|
||||
value=document_ids,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class TestCrud:
|
||||
def test_add_then_query_returns_node(self, store) -> None:
|
||||
node = make_node("n1", 1)
|
||||
@@ -267,8 +281,23 @@ class TestCrud:
|
||||
|
||||
|
||||
class TestBuildWhere:
|
||||
def test_ne_filter_translates_to_not_equal_clause(self) -> None:
|
||||
where, params = _build_where(_ne_filter(1))
|
||||
@pytest.fixture
|
||||
def conn(self) -> Generator[sqlite3.Connection, None, None]:
|
||||
"""A bare connection, sufficient for _build_where(): it only ever
|
||||
touches the connection via PermittedIdsTable, which needs no vec0
|
||||
extension loaded.
|
||||
"""
|
||||
connection = sqlite3.connect(":memory:")
|
||||
try:
|
||||
yield connection
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def test_ne_filter_translates_to_not_equal_clause(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
) -> None:
|
||||
where, params = _build_where(conn, _ne_filter(1))
|
||||
assert where == "(document_id != ?)"
|
||||
assert params == [1]
|
||||
|
||||
@@ -280,7 +309,60 @@ class TestBuildWhere:
|
||||
"b1",
|
||||
]
|
||||
|
||||
def test_fails_closed_when_no_filter_is_translatable(self) -> None:
|
||||
def test_nin_filter_translates_to_not_in_clause(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
) -> None:
|
||||
where, params = _build_where(conn, _nin_filter([1, 2]))
|
||||
assert where == "(document_id NOT IN (?,?))"
|
||||
assert params == [1, 2]
|
||||
|
||||
def test_query_with_nin_filter_excludes_matching_documents(self, store) -> None:
|
||||
store.add([make_node("a1", 1), make_node("b1", 2), make_node("c1", 3)])
|
||||
assert sorted(
|
||||
_query(store, [0.0] * DIM, top_k=5, filters=_nin_filter([1, 2])).ids,
|
||||
) == ["c1"]
|
||||
|
||||
def test_empty_in_filter_excludes_everything(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- An IN filter with an empty value list
|
||||
WHEN:
|
||||
- _build_where() translates it to SQL
|
||||
THEN:
|
||||
- It excludes everything (the opposite of an empty NOT IN
|
||||
filter) -- an empty inclusion list must never widen results
|
||||
"""
|
||||
where, params = _build_where(conn, _in_filter([]))
|
||||
assert where == "(1 = 0)"
|
||||
assert params == []
|
||||
|
||||
def test_empty_nin_filter_excludes_nothing(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- A NOT IN filter with an empty value list -- e.g. an
|
||||
unrestricted chat caller when nothing is currently trashed
|
||||
WHEN:
|
||||
- _build_where() translates it to SQL
|
||||
THEN:
|
||||
- It excludes nothing (unlike an empty IN filter, which
|
||||
excludes everything) -- an empty exclusion list must never
|
||||
narrow results
|
||||
"""
|
||||
where, params = _build_where(conn, _nin_filter([]))
|
||||
assert where == "(1 = 1)"
|
||||
assert params == []
|
||||
|
||||
def test_fails_closed_when_no_filter_is_translatable(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
) -> None:
|
||||
# A nested MetadataFilters is not a MetadataFilter, so it is skipped.
|
||||
# With no translatable clauses, the function must fail closed rather
|
||||
# than emit "()" (invalid SQL) and never widen document access.
|
||||
@@ -293,35 +375,88 @@ class TestBuildWhere:
|
||||
),
|
||||
],
|
||||
)
|
||||
where, params = _build_where(MetadataFilters(filters=[nested]))
|
||||
where, params = _build_where(conn, MetadataFilters(filters=[nested]))
|
||||
assert where == "1 = 0"
|
||||
assert params == []
|
||||
|
||||
def test_fails_closed_when_in_filter_exceeds_max_values(
|
||||
@pytest.mark.parametrize(
|
||||
("build_filter", "sql_op"),
|
||||
[(_in_filter, "IN"), (_nin_filter, "NOT IN")],
|
||||
ids=["in", "nin"],
|
||||
)
|
||||
def test_filter_over_max_values_uses_permitted_ids_table(
|
||||
self,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
conn: sqlite3.Connection,
|
||||
build_filter: Callable[[list[str]], MetadataFilters],
|
||||
sql_op: str,
|
||||
) -> 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)
|
||||
- An IN or NOT IN filter with more values than _MAX_IN_VALUES
|
||||
(SQLite's own bound-parameter limit is 32766; this threshold
|
||||
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
|
||||
- It builds a subquery against PermittedIdsTable's TEMP TABLE,
|
||||
loaded with every id, instead of a literal list SQLite would
|
||||
reject past its own limit -- true for NOT IN too (e.g. an
|
||||
install with an enormous trash), not just IN
|
||||
"""
|
||||
oversized = _in_filter([str(i) for i in range(_MAX_IN_VALUES + 1)])
|
||||
ids = list(range(_MAX_IN_VALUES + 1))
|
||||
oversized = build_filter([str(i) for i in ids])
|
||||
|
||||
with caplog.at_level("WARNING"):
|
||||
where, params = _build_where(oversized)
|
||||
where, params = _build_where(conn, oversized)
|
||||
|
||||
assert where == "(1 = 0)"
|
||||
assert where == (
|
||||
f"(document_id {sql_op} (SELECT id FROM {PermittedIdsTable.TABLE_NAME}))"
|
||||
)
|
||||
assert params == []
|
||||
assert "document_id" in caplog.text
|
||||
loaded = [
|
||||
row[0]
|
||||
for row in conn.execute(
|
||||
f"SELECT id FROM {PermittedIdsTable.TABLE_NAME} ORDER BY id",
|
||||
)
|
||||
]
|
||||
assert loaded == ids
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("build_filter", "expected_ids"),
|
||||
[(_in_filter, ["b1", "c1"]), (_nin_filter, ["a1"])],
|
||||
ids=["in", "nin"],
|
||||
)
|
||||
def test_query_and_get_nodes_scope_correctly_when_filter_exceeds_max_values(
|
||||
self,
|
||||
store: PaperlessSqliteVecVectorStore,
|
||||
mocker: MockerFixture,
|
||||
build_filter: Callable[[list[int]], MetadataFilters],
|
||||
expected_ids: list[str],
|
||||
) -> None:
|
||||
"""
|
||||
GIVEN:
|
||||
- _MAX_IN_VALUES lowered so a small IN/NOT IN filter exceeds it
|
||||
WHEN:
|
||||
- query() and get_nodes() are called with that filter
|
||||
THEN:
|
||||
- Both still correctly scope results -- the PermittedIdsTable
|
||||
temp-table path behaves identically to the literal
|
||||
IN(...)/NOT IN(...) path it replaces above the threshold
|
||||
"""
|
||||
mocker.patch("paperless_ai.vector_store._MAX_IN_VALUES", 1)
|
||||
store.add(
|
||||
[
|
||||
make_node("a1", 1, seed=0.0),
|
||||
make_node("b1", 2, seed=1.0),
|
||||
make_node("c1", 3, seed=2.0),
|
||||
],
|
||||
)
|
||||
|
||||
result = _query(store, [0.0] * DIM, top_k=10, filters=build_filter([2, 3]))
|
||||
nodes = store.get_nodes(filters=build_filter([2, 3]))
|
||||
|
||||
assert sorted(result.ids) == expected_ids
|
||||
assert sorted(n.node_id for n in nodes) == expected_ids
|
||||
|
||||
def test_query_with_untranslatable_filter_returns_no_rows(
|
||||
self,
|
||||
|
||||
@@ -30,6 +30,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.tables import PermittedIdsTable
|
||||
|
||||
logger = logging.getLogger("paperless_ai.vector_store")
|
||||
|
||||
@@ -75,14 +76,12 @@ 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.
|
||||
# _build_where(): the largest IN value list translated into a literal
|
||||
# IN (?,?,...) clause. 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. Above this threshold _build_where() switches to
|
||||
# PermittedIdsTable instead of failing closed -- see its docstring.
|
||||
_MAX_IN_VALUES = 32700
|
||||
|
||||
|
||||
@@ -106,13 +105,21 @@ def _vec0_params(rows: list[_Row]) -> list[tuple[str, int, str, bytes]]:
|
||||
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.
|
||||
def _build_where(
|
||||
conn: sqlite3.Connection,
|
||||
filters: MetadataFilters | None,
|
||||
) -> tuple[str, list[int]]:
|
||||
"""Translate the EQ / IN / NIN / 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.
|
||||
|
||||
``conn`` is only used for an IN/NOT IN filter over _MAX_IN_VALUES: it
|
||||
loads the ids into PermittedIdsTable's TEMP TABLE on that connection
|
||||
rather than binding them as SQL parameters.
|
||||
"""
|
||||
if filters is None or not filters.filters:
|
||||
return "", []
|
||||
@@ -125,28 +132,28 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]:
|
||||
continue
|
||||
if f.key not in _FILTER_COLUMNS: # pragma: no cover - we build the keys
|
||||
raise NotImplementedError(f"Unsupported filter column: {f.key}")
|
||||
if f.operator == FilterOperator.IN:
|
||||
if f.operator in (FilterOperator.IN, FilterOperator.NIN):
|
||||
is_in = f.operator == FilterOperator.IN
|
||||
sql_op = "IN" if is_in else "NOT IN"
|
||||
values = [int(v) for v in f.value] # type: ignore[union-attr]
|
||||
if not values: # pragma: no cover
|
||||
clauses.append("1 = 0")
|
||||
if not values:
|
||||
# An empty IN list matches nothing; an empty NOT IN list
|
||||
# excludes nothing, so it matches everything.
|
||||
clauses.append("1 = 0" if is_in else "1 = 1")
|
||||
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,
|
||||
# A literal list this large would exceed SQLite's own
|
||||
# bound-parameter limit. Load the ids into a TEMP TABLE on
|
||||
# this connection instead and filter via subquery, which has
|
||||
# no such limit -- see PermittedIdsTable. Applies to NOT IN
|
||||
# too (e.g. an install with an enormous trash), not just IN.
|
||||
PermittedIdsTable.load(conn, values)
|
||||
clauses.append(
|
||||
f"{f.key} {sql_op} (SELECT id FROM {PermittedIdsTable.TABLE_NAME})",
|
||||
)
|
||||
clauses.append("1 = 0")
|
||||
continue
|
||||
placeholders = ",".join("?" for _ in values)
|
||||
clauses.append(f"{f.key} IN ({placeholders})")
|
||||
clauses.append(f"{f.key} {sql_op} ({placeholders})")
|
||||
params.extend(values)
|
||||
elif f.operator == FilterOperator.EQ:
|
||||
clauses.append(f"{f.key} = ?")
|
||||
@@ -154,7 +161,7 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]:
|
||||
elif f.operator == FilterOperator.NE:
|
||||
clauses.append(f"{f.key} != ?")
|
||||
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/NIN/NE filters
|
||||
raise NotImplementedError(f"Unsupported filter operator: {f.operator}")
|
||||
if not clauses:
|
||||
# Filters were requested but none could be translated. Fail closed
|
||||
@@ -482,7 +489,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
||||
)
|
||||
if not self.table_exists():
|
||||
return []
|
||||
where, params = _build_where(filters)
|
||||
where, params = _build_where(self._conn, filters)
|
||||
sql = "SELECT node_content, embedding FROM " + DEFAULT_TABLE_NAME
|
||||
if where:
|
||||
sql += " WHERE " + where
|
||||
@@ -498,7 +505,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
||||
if query.query_embedding is None: # pragma: no cover
|
||||
return VectorStoreQueryResult(nodes=[], similarities=[], ids=[])
|
||||
top_k = query.similarity_top_k if query.similarity_top_k is not None else 10
|
||||
where, params = _build_where(query.filters)
|
||||
where, params = _build_where(self._conn, query.filters)
|
||||
sql = (
|
||||
"SELECT id, node_content, embedding, distance FROM "
|
||||
+ DEFAULT_TABLE_NAME
|
||||
|
||||
Reference in New Issue
Block a user