mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-09-02 16:07:15 +00:00
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7c0d3c2f86 | ||
|
|
ec7745e71f | ||
|
|
98e87d91ad |
@@ -1,11 +1,18 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
from unittest import mock
|
from unittest import mock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from django.contrib.auth.models import Permission
|
||||||
from django.contrib.auth.models import User
|
from django.contrib.auth.models import User
|
||||||
from rest_framework import status
|
from rest_framework import status
|
||||||
|
from rest_framework.test import APIClient
|
||||||
from rest_framework.test import APITestCase
|
from rest_framework.test import APITestCase
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from pytest_mock import MockerFixture
|
||||||
|
|
||||||
|
|
||||||
class TestChatStreamingViewInputValidation(APITestCase):
|
class TestChatStreamingViewInputValidation(APITestCase):
|
||||||
def setUp(self) -> None:
|
def setUp(self) -> None:
|
||||||
@@ -42,3 +49,70 @@ class TestChatStreamingViewInputValidation(APITestCase):
|
|||||||
format="json",
|
format="json",
|
||||||
)
|
)
|
||||||
assert resp.status_code == status.HTTP_400_BAD_REQUEST
|
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
|
||||||
|
a caller who can see every document, i.e. an active superuser.
|
||||||
|
"""
|
||||||
|
|
||||||
|
@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, additionally granted
|
||||||
|
view_document -- able to see every document without being a
|
||||||
|
superuser.
|
||||||
|
"""
|
||||||
|
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 view_document
|
||||||
|
WHEN:
|
||||||
|
- They post a chat question with no document_id
|
||||||
|
THEN:
|
||||||
|
- stream_chat_with_documents is called with unrestricted=True for
|
||||||
|
the superuser and unrestricted=False for the regular user, even
|
||||||
|
though that user can view every document
|
||||||
|
"""
|
||||||
|
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_document_ids
|
||||||
from documents.permissions import permitted_object_ids
|
from documents.permissions import permitted_object_ids
|
||||||
from documents.permissions import set_permissions_for_object
|
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.plugins.date_parsing import get_date_parser
|
||||||
from documents.schema import generate_object_with_permissions_schema
|
from documents.schema import generate_object_with_permissions_schema
|
||||||
from documents.search import SearchHit
|
from documents.search import SearchHit
|
||||||
@@ -2329,10 +2330,12 @@ class ChatStreamingView(GenericAPIView[Any]):
|
|||||||
return HttpResponseForbidden("Insufficient permissions")
|
return HttpResponseForbidden("Insufficient permissions")
|
||||||
|
|
||||||
documents = Document.objects.filter(pk=document.pk)
|
documents = Document.objects.filter(pk=document.pk)
|
||||||
|
unrestricted = False
|
||||||
else:
|
else:
|
||||||
documents = Document.objects.filter(
|
documents = Document.objects.filter(
|
||||||
id__in=permitted_document_ids(request.user),
|
id__in=permitted_document_ids(request.user),
|
||||||
)
|
)
|
||||||
|
unrestricted = user_is_unrestricted(request.user)
|
||||||
|
|
||||||
output_language = get_llm_output_language(
|
output_language = get_llm_output_language(
|
||||||
ai_config=ai_config,
|
ai_config=ai_config,
|
||||||
@@ -2343,6 +2346,7 @@ class ChatStreamingView(GenericAPIView[Any]):
|
|||||||
stream_chat_with_documents(
|
stream_chat_with_documents(
|
||||||
query_str=question,
|
query_str=question,
|
||||||
documents=documents,
|
documents=documents,
|
||||||
|
unrestricted=unrestricted,
|
||||||
output_language=output_language,
|
output_language=output_language,
|
||||||
),
|
),
|
||||||
content_type="text/event-stream",
|
content_type="text/event-stream",
|
||||||
|
|||||||
@@ -95,12 +95,15 @@ def _format_chat_metadata_trailer(references: list[dict[str, int | str]]) -> str
|
|||||||
def stream_chat_with_documents(
|
def stream_chat_with_documents(
|
||||||
query_str: str,
|
query_str: str,
|
||||||
documents: QuerySet[Document],
|
documents: QuerySet[Document],
|
||||||
|
*,
|
||||||
|
unrestricted: bool = False,
|
||||||
output_language: str | None = None,
|
output_language: str | None = None,
|
||||||
):
|
):
|
||||||
try:
|
try:
|
||||||
yield from _stream_chat_with_documents(
|
yield from _stream_chat_with_documents(
|
||||||
query_str,
|
query_str,
|
||||||
documents,
|
documents,
|
||||||
|
unrestricted=unrestricted,
|
||||||
output_language=output_language,
|
output_language=output_language,
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -111,6 +114,8 @@ def stream_chat_with_documents(
|
|||||||
def _stream_chat_with_documents(
|
def _stream_chat_with_documents(
|
||||||
query_str: str,
|
query_str: str,
|
||||||
documents: QuerySet[Document],
|
documents: QuerySet[Document],
|
||||||
|
*,
|
||||||
|
unrestricted: bool = False,
|
||||||
output_language: str | None = None,
|
output_language: str | None = None,
|
||||||
):
|
):
|
||||||
if not documents.exists():
|
if not documents.exists():
|
||||||
@@ -123,9 +128,15 @@ def _stream_chat_with_documents(
|
|||||||
from llama_index.core.retrievers import VectorIndexRetriever
|
from llama_index.core.retrievers import VectorIndexRetriever
|
||||||
|
|
||||||
config = AIConfig()
|
config = AIConfig()
|
||||||
filters = _document_id_filters(
|
if unrestricted:
|
||||||
str(pk) for pk in documents.values_list("pk", flat=True)
|
# The caller can see every document, so an id filter would never narrow
|
||||||
)
|
# the search, only risk exceeding the vector store's bound parameter
|
||||||
|
# limit (_MAX_IN_VALUES in vector_store.py) on large installs.
|
||||||
|
filters = None
|
||||||
|
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
|
# Hold the shared read lock for the whole operation: the query engine
|
||||||
# retrieves from the vector store again during synthesis, so the connection
|
# retrieves from the vector store again during synthesis, so the connection
|
||||||
|
|||||||
@@ -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:
|
class IndexMetaTable:
|
||||||
"""Typed accessors over index_meta's key/value rows -- replaces
|
"""Typed accessors over index_meta's key/value rows -- replaces
|
||||||
PaperlessSqliteVecVectorStore._meta_get_on/_meta_set_on, which returned
|
PaperlessSqliteVecVectorStore._meta_get_on/_meta_set_on, which returned
|
||||||
|
|||||||
@@ -1,4 +1,8 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
from typing import Any
|
||||||
from unittest.mock import MagicMock
|
from unittest.mock import MagicMock
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|
||||||
@@ -18,6 +22,11 @@ from paperless_ai.chat import _build_chat_prompt
|
|||||||
from paperless_ai.chat import _build_refine_prompt
|
from paperless_ai.chat import _build_refine_prompt
|
||||||
from paperless_ai.chat import stream_chat_with_documents
|
from paperless_ai.chat import stream_chat_with_documents
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest_mock
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(autouse=True)
|
@pytest.fixture(autouse=True)
|
||||||
def patch_embed_model():
|
def patch_embed_model():
|
||||||
@@ -312,6 +321,30 @@ def test_stream_chat_unexpected_failure_returns_generic_error(caplog) -> None:
|
|||||||
|
|
||||||
@pytest.mark.django_db
|
@pytest.mark.django_db
|
||||||
class TestStreamChatRetrieval:
|
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(
|
def test_no_nodes_yields_no_content_message(
|
||||||
self,
|
self,
|
||||||
temp_llm_index_dir,
|
temp_llm_index_dir,
|
||||||
@@ -329,9 +362,9 @@ class TestStreamChatRetrieval:
|
|||||||
|
|
||||||
def test_chat_filter_contains_only_requested_document_ids(
|
def test_chat_filter_contains_only_requested_document_ids(
|
||||||
self,
|
self,
|
||||||
temp_llm_index_dir,
|
temp_llm_index_dir: Path,
|
||||||
mock_embed_model,
|
mock_embed_model: pytest_mock.MockType,
|
||||||
mocker,
|
captured_filters: list[Any],
|
||||||
) -> None:
|
) -> None:
|
||||||
"""The MetadataFilter passed to the retriever must be scoped to the
|
"""The MetadataFilter passed to the retriever must be scoped to the
|
||||||
requested documents only — content from other indexed documents must
|
requested documents only — content from other indexed documents must
|
||||||
@@ -342,22 +375,6 @@ class TestStreamChatRetrieval:
|
|||||||
indexing.llm_index_add_or_update_document(included)
|
indexing.llm_index_add_or_update_document(included)
|
||||||
indexing.llm_index_add_or_update_document(excluded)
|
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(
|
list(
|
||||||
chat.stream_chat_with_documents(
|
chat.stream_chat_with_documents(
|
||||||
"question?",
|
"question?",
|
||||||
@@ -372,6 +389,36 @@ class TestStreamChatRetrieval:
|
|||||||
assert str(included.pk) in filter_values
|
assert str(included.pk) in filter_values
|
||||||
assert str(excluded.pk) not in filter_values
|
assert str(excluded.pk) not in filter_values
|
||||||
|
|
||||||
|
def test_unrestricted_chat_skips_document_id_filter(
|
||||||
|
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
|
||||||
|
WHEN:
|
||||||
|
- stream_chat_with_documents is called with unrestricted=True
|
||||||
|
THEN:
|
||||||
|
- The retriever receives no document id filter (filters=None), so
|
||||||
|
the whole index is searched instead of an IN-list that 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 captured_filters, "VectorIndexRetriever was never constructed"
|
||||||
|
assert captured_filters[0] is None
|
||||||
|
|
||||||
@pytest.mark.django_db
|
@pytest.mark.django_db
|
||||||
def test_get_document_references_only_queries_referenced_documents(
|
def test_get_document_references_only_queries_referenced_documents(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ from paperless_ai.tables import DocumentChunksTable
|
|||||||
from paperless_ai.tables import DocumentMetaRow
|
from paperless_ai.tables import DocumentMetaRow
|
||||||
from paperless_ai.tables import DocumentMetaTable
|
from paperless_ai.tables import DocumentMetaTable
|
||||||
from paperless_ai.tables import IndexMetaTable
|
from paperless_ai.tables import IndexMetaTable
|
||||||
|
from paperless_ai.tables import PermittedIdsTable
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
@@ -338,3 +339,85 @@ class TestIndexMetaTable:
|
|||||||
IndexMetaTable.increment_total_inserts(conn, 100)
|
IndexMetaTable.increment_total_inserts(conn, 100)
|
||||||
IndexMetaTable.reset_total_inserts(conn, 7)
|
IndexMetaTable.reset_total_inserts(conn, 7)
|
||||||
assert IndexMetaTable.get_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
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ from paperless_ai.migrations import Migration
|
|||||||
from paperless_ai.migrations import m0001_v1_to_v2
|
from paperless_ai.migrations import m0001_v1_to_v2
|
||||||
from paperless_ai.tables import DocumentChunksTable
|
from paperless_ai.tables import DocumentChunksTable
|
||||||
from paperless_ai.tables import DocumentMetaTable
|
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 _MAX_IN_VALUES
|
||||||
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
|
||||||
@@ -267,8 +268,23 @@ class TestCrud:
|
|||||||
|
|
||||||
|
|
||||||
class TestBuildWhere:
|
class TestBuildWhere:
|
||||||
def test_ne_filter_translates_to_not_equal_clause(self) -> None:
|
@pytest.fixture
|
||||||
where, params = _build_where(_ne_filter(1))
|
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 where == "(document_id != ?)"
|
||||||
assert params == [1]
|
assert params == [1]
|
||||||
|
|
||||||
@@ -280,7 +296,10 @@ class TestBuildWhere:
|
|||||||
"b1",
|
"b1",
|
||||||
]
|
]
|
||||||
|
|
||||||
def test_fails_closed_when_no_filter_is_translatable(self) -> None:
|
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.
|
# A nested MetadataFilters is not a MetadataFilter, so it is skipped.
|
||||||
# With no translatable clauses, the function must fail closed rather
|
# With no translatable clauses, the function must fail closed rather
|
||||||
# than emit "()" (invalid SQL) and never widen document access.
|
# than emit "()" (invalid SQL) and never widen document access.
|
||||||
@@ -293,35 +312,74 @@ class TestBuildWhere:
|
|||||||
),
|
),
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
where, params = _build_where(MetadataFilters(filters=[nested]))
|
where, params = _build_where(conn, MetadataFilters(filters=[nested]))
|
||||||
assert where == "1 = 0"
|
assert where == "1 = 0"
|
||||||
assert params == []
|
assert params == []
|
||||||
|
|
||||||
def test_fails_closed_when_in_filter_exceeds_max_values(
|
def test_in_filter_over_max_values_uses_permitted_ids_table(
|
||||||
self,
|
self,
|
||||||
caplog: pytest.LogCaptureFixture,
|
conn: sqlite3.Connection,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""
|
||||||
GIVEN:
|
GIVEN:
|
||||||
- An IN filter with more values than _MAX_IN_VALUES (SQLite's
|
- An IN filter with more values than _MAX_IN_VALUES (SQLite's
|
||||||
own bound-parameter limit is 32766; this guard sits below
|
own bound-parameter limit is 32766; this threshold sits below
|
||||||
that with headroom for the query's other bound parameters)
|
that with headroom for the query's other bound parameters)
|
||||||
WHEN:
|
WHEN:
|
||||||
- _build_where() translates it to SQL
|
- _build_where() translates it to SQL
|
||||||
THEN:
|
THEN:
|
||||||
- It fails closed ("1 = 0", no params) instead of building an
|
- It builds a subquery against PermittedIdsTable's TEMP TABLE,
|
||||||
IN clause SQLite would reject, and logs a warning -- this
|
loaded with every id, instead of a literal IN(...) list that
|
||||||
filter scopes document access, so refusing to build it must
|
SQLite would reject past its own limit -- the filter still
|
||||||
never widen the scope to "everything" by accident
|
scopes document access to exactly the requested ids, never
|
||||||
|
widening the scope to "everything"
|
||||||
"""
|
"""
|
||||||
oversized = _in_filter([str(i) for i in range(_MAX_IN_VALUES + 1)])
|
ids = list(range(_MAX_IN_VALUES + 1))
|
||||||
|
oversized = _in_filter([str(i) for i in ids])
|
||||||
|
|
||||||
with caplog.at_level("WARNING"):
|
where, params = _build_where(conn, oversized)
|
||||||
where, params = _build_where(oversized)
|
|
||||||
|
|
||||||
assert where == "(1 = 0)"
|
assert where == (
|
||||||
|
f"(document_id IN (SELECT id FROM {PermittedIdsTable.TABLE_NAME}))"
|
||||||
|
)
|
||||||
assert params == []
|
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
|
||||||
|
|
||||||
|
def test_query_and_get_nodes_scope_correctly_when_in_filter_exceeds_max_values(
|
||||||
|
self,
|
||||||
|
store: PaperlessSqliteVecVectorStore,
|
||||||
|
mocker: MockerFixture,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
GIVEN:
|
||||||
|
- _MAX_IN_VALUES lowered so a small IN filter exceeds it
|
||||||
|
WHEN:
|
||||||
|
- query() and get_nodes() are called with that filter
|
||||||
|
THEN:
|
||||||
|
- Both still correctly scope results to the permitted ids -- the
|
||||||
|
PermittedIdsTable temp-table path behaves identically to the
|
||||||
|
literal 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=_in_filter([2, 3]))
|
||||||
|
nodes = store.get_nodes(filters=_in_filter([2, 3]))
|
||||||
|
|
||||||
|
assert sorted(result.ids) == ["b1", "c1"]
|
||||||
|
assert sorted(n.node_id for n in nodes) == ["b1", "c1"]
|
||||||
|
|
||||||
def test_query_with_untranslatable_filter_returns_no_rows(
|
def test_query_with_untranslatable_filter_returns_no_rows(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -30,6 +30,7 @@ from paperless_ai.tables import DocumentChunksTable
|
|||||||
from paperless_ai.tables import DocumentMetaRow
|
from paperless_ai.tables import DocumentMetaRow
|
||||||
from paperless_ai.tables import DocumentMetaTable
|
from paperless_ai.tables import DocumentMetaTable
|
||||||
from paperless_ai.tables import IndexMetaTable
|
from paperless_ai.tables import IndexMetaTable
|
||||||
|
from paperless_ai.tables import PermittedIdsTable
|
||||||
|
|
||||||
logger = logging.getLogger("paperless_ai.vector_store")
|
logger = logging.getLogger("paperless_ai.vector_store")
|
||||||
|
|
||||||
@@ -75,14 +76,12 @@ class _Row(NamedTuple):
|
|||||||
embedding: bytes
|
embedding: bytes
|
||||||
|
|
||||||
|
|
||||||
# _build_where(): the largest IN value list translated into bound SQL
|
# _build_where(): the largest IN value list translated into a literal
|
||||||
# parameters. SQLite's own hard limit (SQLITE_MAX_VARIABLE_NUMBER) is 32766
|
# IN (?,?,...) clause. SQLite's own hard limit (SQLITE_MAX_VARIABLE_NUMBER)
|
||||||
# by default; this leaves headroom below that for the query's other bound
|
# is 32766 by default; this leaves headroom below that for the query's other
|
||||||
# parameters (the embedding blob, k, and any NE clause) and for the limit
|
# bound parameters (the embedding blob, k, and any NE clause) and for the
|
||||||
# itself to move. An IN filter this large should not happen in practice --
|
# limit itself to move. Above this threshold _build_where() switches to
|
||||||
# callers are expected to pass None (no filter) rather than every id when
|
# PermittedIdsTable instead of failing closed -- see its docstring.
|
||||||
# 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
|
_MAX_IN_VALUES = 32700
|
||||||
|
|
||||||
|
|
||||||
@@ -106,13 +105,20 @@ 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]
|
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]]:
|
def _build_where(
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
filters: MetadataFilters | None,
|
||||||
|
) -> tuple[str, list[int]]:
|
||||||
"""Translate the EQ / IN / NE filters we use into a parameterized SQL
|
"""Translate the EQ / IN / NE filters we use into a parameterized SQL
|
||||||
clause on vec0 metadata columns. Returns ("", []) when there is nothing
|
clause on vec0 metadata columns. Returns ("", []) when there is nothing
|
||||||
to filter. document_id is vec0's only filterable column and is INTEGER;
|
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
|
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
|
strings in places, e.g. indexing.py's MetadataFilter construction) don't
|
||||||
have to be individually correct -- vec0 doesn't coerce types itself.
|
have to be individually correct -- vec0 doesn't coerce types itself.
|
||||||
|
|
||||||
|
``conn`` is only used for an 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:
|
if filters is None or not filters.filters:
|
||||||
return "", []
|
return "", []
|
||||||
@@ -131,19 +137,14 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]:
|
|||||||
clauses.append("1 = 0")
|
clauses.append("1 = 0")
|
||||||
continue
|
continue
|
||||||
if len(values) > _MAX_IN_VALUES:
|
if len(values) > _MAX_IN_VALUES:
|
||||||
# Fail closed (see the empty-clauses case below) rather than
|
# A literal IN(...) list this large would exceed SQLite's own
|
||||||
# let SQLite raise "too many SQL variables" past its own
|
# bound-parameter limit. Load the ids into a TEMP TABLE on
|
||||||
# limit: this filter scopes document access, so an IN list
|
# this connection instead and filter via subquery, which has
|
||||||
# too large to safely bind must match no rows, never widen
|
# no such limit -- see PermittedIdsTable.
|
||||||
# the scope to "everything" by accident.
|
PermittedIdsTable.load(conn, values)
|
||||||
logger.warning(
|
clauses.append(
|
||||||
"Refusing to build an IN filter on %r with %d values "
|
f"{f.key} IN (SELECT id FROM {PermittedIdsTable.TABLE_NAME})",
|
||||||
"(over the %d-value safety limit); returning no rows.",
|
|
||||||
f.key,
|
|
||||||
len(values),
|
|
||||||
_MAX_IN_VALUES,
|
|
||||||
)
|
)
|
||||||
clauses.append("1 = 0")
|
|
||||||
continue
|
continue
|
||||||
placeholders = ",".join("?" for _ in values)
|
placeholders = ",".join("?" for _ in values)
|
||||||
clauses.append(f"{f.key} IN ({placeholders})")
|
clauses.append(f"{f.key} IN ({placeholders})")
|
||||||
@@ -482,7 +483,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
)
|
)
|
||||||
if not self.table_exists():
|
if not self.table_exists():
|
||||||
return []
|
return []
|
||||||
where, params = _build_where(filters)
|
where, params = _build_where(self._conn, filters)
|
||||||
sql = "SELECT node_content, embedding FROM " + DEFAULT_TABLE_NAME
|
sql = "SELECT node_content, embedding FROM " + DEFAULT_TABLE_NAME
|
||||||
if where:
|
if where:
|
||||||
sql += " WHERE " + where
|
sql += " WHERE " + where
|
||||||
@@ -498,7 +499,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
|
|||||||
if query.query_embedding is None: # pragma: no cover
|
if query.query_embedding is None: # pragma: no cover
|
||||||
return VectorStoreQueryResult(nodes=[], similarities=[], ids=[])
|
return VectorStoreQueryResult(nodes=[], similarities=[], ids=[])
|
||||||
top_k = query.similarity_top_k if query.similarity_top_k is not None else 10
|
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 = (
|
sql = (
|
||||||
"SELECT id, node_content, embedding, distance FROM "
|
"SELECT id, node_content, embedding, distance FROM "
|
||||||
+ DEFAULT_TABLE_NAME
|
+ DEFAULT_TABLE_NAME
|
||||||
|
|||||||
Reference in New Issue
Block a user