From 62d729e30146e2e866ac81957e59c031bde60b84 Mon Sep 17 00:00:00 2001 From: Trenton H <797416+stumpylog@users.noreply.github.com> Date: Thu, 12 Nov 2026 12:08:59 -0800 Subject: [PATCH] Fix: Allow AI chat filters over more than 32700 documents (#14754) --- src/paperless_ai/chat.py | 34 +++++++--- src/paperless_ai/tables.py | 9 +++ src/paperless_ai/tests/test_chat.py | 71 +++++++++++++++++++-- src/paperless_ai/tests/test_vector_store.py | 44 +++++++------ src/paperless_ai/vector_store.py | 41 +++++------- 5 files changed, 141 insertions(+), 58 deletions(-) diff --git a/src/paperless_ai/chat.py b/src/paperless_ai/chat.py index 088a5b285f..bf1c720bc1 100644 --- a/src/paperless_ai/chat.py +++ b/src/paperless_ai/chat.py @@ -46,6 +46,16 @@ def _build_document_reference( } +def _permitted_documents_filter(store, permitted_ids: set[int]): + """Scope retrieval to ``permitted_ids`` using whichever of an inclusion or + exclusion list is shorter, since a long inclusion list is much slower for + the vector store than a long exclusion list.""" + denied_ids = store.indexed_document_ids() - permitted_ids + if len(permitted_ids) <= len(denied_ids): + return document_id_filters(str(pk) for pk in permitted_ids) + return exclude_document_ids_filter(str(pk) for pk in denied_ids) + + def _get_document_references( documents: QuerySet[Document], top_nodes: list, @@ -128,24 +138,23 @@ def _stream_chat_with_documents( from llama_index.core.retrievers import VectorIndexRetriever config = AIConfig() + permitted_ids: set[int] | None = None 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. + # Exclude trashed ids (usually few) rather than send every permitted + # id on each query. Trashed documents stay indexed until permanent + # deletion (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) - ) + permitted_ids = set(documents.values_list("pk", flat=True)) # Hold the shared read lock only while reading the index, so the compaction # swap cannot replace the database file while this connection is open. # Everything after retrieval works on the nodes already in memory, so a slow # LLM stream never holds the lock. with read_store() as store: + if permitted_ids is not None: + filters = _permitted_documents_filter(store, permitted_ids) index = load_or_build_index(config, store) retriever = VectorIndexRetriever( index=index, @@ -159,6 +168,15 @@ def _stream_chat_with_documents( with db_connection_released(): top_nodes = retriever.retrieve(query_str) + if permitted_ids is not None: + # An exclusion filter fails open for a document indexed after the + # permitted set was read, so drop anything not explicitly permitted. + top_nodes = [ + node + for node in top_nodes + if int(node.metadata["document_id"]) in permitted_ids + ] + if not top_nodes: logger.warning("No nodes found for the given documents.") yield CHAT_NO_CONTENT_MESSAGE diff --git a/src/paperless_ai/tables.py b/src/paperless_ai/tables.py index c8d6d838c0..621a73f03b 100644 --- a/src/paperless_ai/tables.py +++ b/src/paperless_ai/tables.py @@ -141,6 +141,15 @@ class DocumentMetaTable: (DocumentMetaRow(r["document_id"], r["modified"]) for r in batch), ) + @staticmethod + def all_document_ids(conn: sqlite3.Connection) -> set[int]: + """Every indexed document_id. document_meta has one row per document, + so this is a cheap primary-key scan, unlike reading the vec0 table.""" + return { + row["document_id"] + for row in conn.execute("SELECT document_id FROM document_meta") + } + @staticmethod def all_modified_times(conn: sqlite3.Connection) -> dict[str, str]: """Full document_id -> modified map, for get_modified_times()'s diff --git a/src/paperless_ai/tests/test_chat.py b/src/paperless_ai/tests/test_chat.py index d92456ed58..358dec8d22 100644 --- a/src/paperless_ai/tests/test_chat.py +++ b/src/paperless_ai/tests/test_chat.py @@ -13,6 +13,7 @@ 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 +from llama_index.core.vector_stores.types import FilterOperator from documents.models import Document from paperless_ai import chat @@ -412,10 +413,8 @@ class TestStreamChatRetrieval: 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 + - The retriever receives a NOT IN filter excluding zero ids + rather than an IN list of every permitted id """ document = DocumentFactory.create(content="indexed document content") indexing.llm_index_add_or_update_document(document) @@ -468,6 +467,70 @@ class TestStreamChatRetrieval: assert str(trashed.pk) in filter_values assert str(kept.pk) not in filter_values + def test_restricted_chat_excludes_denied_ids_when_few_are_denied( + self, + temp_llm_index_dir: Path, + mock_embed_model: pytest_mock.MockType, + captured_filters: list[Any], + ) -> None: + """ + GIVEN: + - Three indexed documents, one of which the caller may not see + WHEN: + - A restricted caller chats + THEN: + - The retriever receives a NOT IN filter on the denied id, not an + IN list of the permitted ids + """ + docs = DocumentFactory.create_batch(3, content="indexed document content") + for doc in docs: + indexing.llm_index_add_or_update_document(doc) + permitted = Document.objects.filter(pk__in=[docs[0].pk, docs[1].pk]) + + list(chat.stream_chat_with_documents("question?", permitted)) + + metadata_filter = captured_filters[0].filters[0] + assert metadata_filter.operator == FilterOperator.NIN + assert metadata_filter.value == [str(docs[2].pk)] + + def test_restricted_chat_drops_retrieved_nodes_of_unpermitted_documents( + self, + temp_llm_index_dir: Path, + mock_embed_model: pytest_mock.MockType, + mocker: pytest_mock.MockerFixture, + ) -> None: + """ + GIVEN: + - A retriever that returns a node for a document the caller may + not see, as when it is indexed after the permitted set was read + WHEN: + - A restricted caller chats + THEN: + - Only the node of the permitted document reaches the synthesizer + """ + permitted_doc = DocumentFactory.create(content="permitted content") + denied_doc = DocumentFactory.create(content="denied content") + permitted_node = MagicMock(metadata={"document_id": str(permitted_doc.pk)}) + denied_node = MagicMock(metadata={"document_id": str(denied_doc.pk)}) + mocker.patch("paperless_ai.chat.AIClient") + retriever = mocker.patch( + "llama_index.core.retrievers.VectorIndexRetriever", + ).return_value + retriever.retrieve.return_value = [denied_node, permitted_node] + synthesizer = mocker.patch( + "llama_index.core.response_synthesizers.get_response_synthesizer", + ).return_value + synthesizer.synthesize.return_value.response_gen = iter(["chunk"]) + + list( + chat.stream_chat_with_documents( + "question?", + Document.objects.filter(pk=permitted_doc.pk), + ), + ) + + synthesizer.synthesize.assert_called_once_with("question?", [permitted_node]) + @pytest.mark.django_db def test_get_document_references_only_queries_referenced_documents( self, diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index 9981b89a74..1fb47e307b 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -18,7 +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.vector_store import _MAX_IN_VALUES +from paperless_ai.vector_store import _MAX_INLINE_IN_VALUES from paperless_ai.vector_store import DB_FILENAME from paperless_ai.vector_store import DB_PAGE_SIZE from paperless_ai.vector_store import DEFAULT_TABLE_NAME @@ -354,37 +354,39 @@ class TestBuildWhere: assert params == [] @pytest.mark.parametrize( - "build_filter", - [_in_filter, _nin_filter], - ids=["in", "nin"], + ("build_filter", "expected_ids"), + [ + pytest.param(_in_filter, ["a", "c"], id="in"), + pytest.param(_nin_filter, ["b"], id="nin"), + ], ) - def test_fails_closed_when_filter_exceeds_max_values( + def test_query_with_oversized_filter_scopes_results( self, + store: PaperlessSqliteVecVectorStore, build_filter: Callable[[list[str]], MetadataFilters], - caplog: pytest.LogCaptureFixture, + expected_ids: list[str], ) -> None: """ GIVEN: - - 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) + - A store holding documents 1, 2 and 40000 + - An IN or NOT IN filter naming documents 1 and 40000, padded + past _MAX_INLINE_IN_VALUES with unindexed ids WHEN: - - _build_where() translates it to SQL + - The store is queried and get_nodes() is called with that filter THEN: - - 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 + - Both return exactly the documents the filter allows """ - oversized = build_filter([str(i) for i in range(_MAX_IN_VALUES + 1)]) + store.add( + [make_node("a", 1), make_node("b", 2), make_node("c", 40000)], + ) + padding = range(100, 100 + _MAX_INLINE_IN_VALUES) + filters = build_filter([str(i) for i in (1, 40000, *padding)]) - with caplog.at_level("WARNING"): - where, params = _build_where(oversized) + result = _query(store, [0.0] * DIM, top_k=5, filters=filters) + nodes = store.get_nodes(filters=filters) - assert where == "(1 = 0)" - assert params == [] - assert "document_id" in caplog.text + 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, diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index 7b48934247..e454be9dcc 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -89,15 +89,10 @@ class _Row(NamedTuple): embedding: bytes -# _build_where(): the largest IN value list translated into bound SQL -# parameters. SQLite's own hard limit (SQLITE_MAX_VARIABLE_NUMBER) is 32766 -# by default; this leaves headroom below that for the query's other bound -# parameters (the embedding blob, k, and any NE clause) and for the limit -# itself to move. An IN filter this large should not happen in practice -- -# callers are expected to pass None (no filter) rather than every id when -# the filter would not actually narrow anything -- so this is a guard -# against a future regression, not a normal code path. -_MAX_IN_VALUES = 32700 +# _build_where(): the largest IN / NOT IN list bound as one parameter per value, +# leaving headroom under SQLite's 32766-variable limit for the query's other +# parameters. Longer lists are bound as a single JSON array via json_each(). +_MAX_INLINE_IN_VALUES = 32700 def _pack(embedding: Sequence[float]) -> bytes: @@ -120,7 +115,7 @@ 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]]: +def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int | str]]: """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 @@ -132,7 +127,7 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: if filters is None or not filters.filters: return "", [] clauses: list[str] = [] - params: list[int] = [] + params: list[int | str] = [] for f in filters.filters: # filters.filters is Union[MetadataFilter, ExactMatchFilter, MetadataFilters]; # we only build MetadataFilter entries, so skip anything else at runtime. @@ -149,21 +144,11 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: # excludes nothing, so it matches everything. clauses.append("1 = 0" if is_in else "1 = 1") continue - if len(values) > _MAX_IN_VALUES: - # 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 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, + if len(values) > _MAX_INLINE_IN_VALUES: + clauses.append( + f"{f.key} {sql_op} (SELECT value FROM json_each(?))", ) - clauses.append("1 = 0") + params.append(json.dumps(values)) continue placeholders = ",".join("?" for _ in values) clauses.append(f"{f.key} {sql_op} ({placeholders})") @@ -558,6 +543,12 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): return {} return DocumentMetaTable.all_modified_times(self._conn) + def indexed_document_ids(self) -> set[int]: + """Return the ids of all indexed documents.""" + if not self.table_exists(): + return set() + return DocumentMetaTable.all_document_ids(self._conn) + @property def _db_path(self) -> str: return str(Path(self._uri) / DB_FILENAME)