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)