Fix: skip vector store document id filter for unrestricted chat users (#13937)

* Fix: skip vector store document id filter for unrestricted chat users

ChatStreamingView built an IN filter from every permitted document id
for the "chat over all documents" case, which exceeds the vector
store's SQLite bound-parameter safety limit on installs with more
than ~32700 documents, silently returning no context. For a user who
can see every document (an active superuser), that filter never
narrows anything, so skip it and let the retriever search the whole
index instead.

* Minor improvements from a Claude review

* When a user is unrestricted chatting, still exclude trashed documents using a 'NOT IN' SQL statement.  Wire that up where we need it

* Update src/paperless_ai/chat.py

Co-authored-by: shamoon <4887959+shamoon@users.noreply.github.com>
This commit is contained in:
Trenton H
2026-09-02 18:08:54 +00:00
committed by GitHub
co-authored by shamoon
parent 5d6ea11828
commit 351892bbab
8 changed files with 327 additions and 56 deletions
+1 -1
View File
@@ -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\
+77
View File
@@ -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
)
+4
View File
@@ -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",
+19 -4
View File
@@ -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,18 @@ 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.deleted_objects.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
+19 -2
View File
@@ -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)
+114 -23
View File
@@ -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,
+70 -9
View File
@@ -1,5 +1,6 @@
import inspect
import sqlite3
from collections.abc import Callable
from collections.abc import Generator
from pathlib import Path
@@ -97,6 +98,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)
@@ -280,6 +293,47 @@ class TestBuildWhere:
"b1",
]
def test_nin_filter_translates_to_not_in_clause(self) -> None:
where, params = _build_where(_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) -> 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(_in_filter([]))
assert where == "(1 = 0)"
assert params == []
def test_empty_nin_filter_excludes_nothing(self) -> 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(_nin_filter([]))
assert where == "(1 = 1)"
assert params == []
def test_fails_closed_when_no_filter_is_translatable(self) -> None:
# A nested MetadataFilters is not a MetadataFilter, so it is skipped.
# With no translatable clauses, the function must fail closed rather
@@ -297,24 +351,31 @@ class TestBuildWhere:
assert where == "1 = 0"
assert params == []
def test_fails_closed_when_in_filter_exceeds_max_values(
@pytest.mark.parametrize(
"build_filter",
[_in_filter, _nin_filter],
ids=["in", "nin"],
)
def test_fails_closed_when_filter_exceeds_max_values(
self,
build_filter: Callable[[list[str]], MetadataFilters],
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)
- An IN or NOT 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
- It fails closed ("1 = 0", no params) instead of building a
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. Failing open on
a NOT IN would surface exactly the excluded rows
"""
oversized = _in_filter([str(i) for i in range(_MAX_IN_VALUES + 1)])
oversized = build_filter([str(i) for i in range(_MAX_IN_VALUES + 1)])
with caplog.at_level("WARNING"):
where, params = _build_where(oversized)
+23 -17
View File
@@ -107,12 +107,13 @@ def _vec0_params(rows: list[_Row]) -> list[tuple[str, int, str, bytes]]:
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.
"""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.
"""
if filters is None or not filters.filters:
return "", []
@@ -125,20 +126,25 @@ 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.
# Refuse rather than risk SQLite's own bound-parameter limit
# ("too many SQL variables"): a list this large must match no
# rows, never widen the scope to "everything" -- true for
# NOT IN too, where failing open would surface every
# excluded row.
logger.warning(
"Refusing to build an IN filter on %r with %d values "
"Refusing to build a %s filter on %r with %d values "
"(over the %d-value safety limit); returning no rows.",
sql_op,
f.key,
len(values),
_MAX_IN_VALUES,
@@ -146,7 +152,7 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]:
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 +160,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