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

@@ -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

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

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,

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)