From 5e4936ee83c235d38426e1a7341bd3b8eba02a4a Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Thu, 12 Nov 2026 08:13:05 -0800 Subject: [PATCH] Fix: Speed up AI chat scoping for users who can see most documents A NOT IN list of the denied ids is as fast as an unfiltered query. Chat now diffs the indexed document ids against the permitted ids and filters with whichever list is shorter. A NOT IN filter would fail open for a document indexed after the permitted ids were read, so retrieved nodes are also checked against the permitted set and dropped when they are not in it. --- src/paperless_ai/chat.py | 26 ++++++++++-- src/paperless_ai/tables.py | 9 ++++ src/paperless_ai/tests/test_chat.py | 65 +++++++++++++++++++++++++++++ src/paperless_ai/vector_store.py | 6 +++ 4 files changed, 103 insertions(+), 3 deletions(-) diff --git a/src/paperless_ai/chat.py b/src/paperless_ai/chat.py index f42081bb27..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,6 +138,7 @@ 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) rather than send every permitted # id on each query. Trashed documents stay indexed until permanent @@ -135,15 +146,15 @@ def _stream_chat_with_documents( 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, @@ -157,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 b6bc3c89e9..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 @@ -466,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/vector_store.py b/src/paperless_ai/vector_store.py index 9e904ce648..e454be9dcc 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -543,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)