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,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
+        # Exclude trashed ids (usually few) rather than send every permitted
-        # full permitted set, which risks the vector store's bound parameter
+        # id on each query. Trashed documents stay indexed until permanent
-        # limit (_MAX_IN_VALUES) on large installs. Trashed documents stay
+        # deletion (post_delete, not trash), so must be excluded explicitly.
-        # 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(
+        permitted_ids = set(documents.values_list("pk", flat=True))
-            str(pk) for pk in 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

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
@@ -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 retriever receives a NOT IN filter excluding zero ids
-              the whole index is effectively searched -- and no IN-list is
+              rather than an IN list of every permitted id
-              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)
@@ -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,

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",
+        ("build_filter", "expected_ids"),
-        [_in_filter, _nin_filter],
+        [
-        ids=["in", "nin"],
+            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
+            - A store holding documents 1, 2 and 40000
-              (SQLite's own bound-parameter limit is 32766; this guard sits
+            - An IN or NOT IN filter naming documents 1 and 40000, padded
-              below that with headroom for the query's other bound parameters)
+              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
+            - Both return exactly the documents the filter allows
-              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 = 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"):
+        result = _query(store, [0.0] * DIM, top_k=5, filters=filters)
-            where, params = _build_where(oversized)
+        nodes = store.get_nodes(filters=filters)
 
-        assert where == "(1 = 0)"
+        assert sorted(result.ids) == expected_ids
-        assert params == []
+        assert sorted(n.node_id for n in nodes) == expected_ids
-        assert "document_id" in caplog.text
 
     def test_query_with_untranslatable_filter_returns_no_rows(
         self,

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
+# _build_where(): the largest IN / NOT IN list bound as one parameter per value,
-# parameters. SQLite's own hard limit (SQLITE_MAX_VARIABLE_NUMBER) is 32766
+# leaving headroom under SQLite's 32766-variable limit for the query's other
-# by default; this leaves headroom below that for the query's other bound
+# parameters. Longer lists are bound as a single JSON array via json_each().
-# parameters (the embedding blob, k, and any NE clause) and for the limit
+_MAX_INLINE_IN_VALUES = 32700
-# 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
 
 
 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:
+            if len(values) > _MAX_INLINE_IN_VALUES:
-                # Refuse rather than risk SQLite's own bound-parameter limit
+                clauses.append(
-                # ("too many SQL variables"): a list this large must match no
+                    f"{f.key} {sql_op} (SELECT value FROM json_each(?))",
-                # 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,
                 )
-                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)