mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-10-14 03:57:13 +00:00
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:
1 parent
5d6ea11828
commit
351892bbab
8 files changed
+327
-56
No files matched your search
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user