diff --git a/pyproject.toml b/pyproject.toml index 361b30740..7ee372782 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -247,7 +247,7 @@ per-file-ignores."src/documents/models.py" = [ isort.force-single-line = true [tool.codespell] -ignore-words-list = "criterias,afterall,valeu,ureue,equest,ure,assertIn,Oktober,commitish" +ignore-words-list = "criterias,afterall,valeu,ureue,equest,ure,assertIn,Oktober,commitish,NIN,nin" skip = """\ src-ui/src/locale/*,src-ui/pnpm-lock.yaml,src-ui/e2e/*,src/paperless_mail/tests/samples/*,src/paperless/tests/samples\ /mail/*,src/documents/tests/samples/*,*.po,*.json\ diff --git a/src/documents/tests/test_api_chat.py b/src/documents/tests/test_api_chat.py index 421c74558..896fa4584 100644 --- a/src/documents/tests/test_api_chat.py +++ b/src/documents/tests/test_api_chat.py @@ -1,11 +1,18 @@ from __future__ import annotations +from typing import TYPE_CHECKING from unittest import mock +import pytest +from django.contrib.auth.models import Permission from django.contrib.auth.models import User from rest_framework import status +from rest_framework.test import APIClient from rest_framework.test import APITestCase +if TYPE_CHECKING: + from pytest_mock import MockerFixture + class TestChatStreamingViewInputValidation(APITestCase): def setUp(self) -> None: @@ -42,3 +49,73 @@ class TestChatStreamingViewInputValidation(APITestCase): format="json", ) assert resp.status_code == status.HTTP_400_BAD_REQUEST + + +@pytest.mark.django_db +class TestChatStreamingViewUnrestrictedFlag: + """The document id filter may only be skipped (``unrestricted=True``) for + an active superuser, never for a regular user -- regardless of what + permissions that user holds. + """ + + @pytest.fixture + def mocked_stream_chat(self, mocker: MockerFixture) -> mock.MagicMock: + """AI enabled, with stream_chat_with_documents patched so the view + never touches the real vector store; returns the patched callable so + tests can inspect how it was called. + """ + mocker.patch("documents.views.AIConfig").return_value.ai_enabled = True + return mocker.patch( + "documents.views.stream_chat_with_documents", + return_value=iter(()), + ) + + @pytest.fixture + def viewer_client(self, user_client: APIClient, regular_user: User) -> APIClient: + """The conftest regular-user client, granted the global + view_document permission -- the minimum ViewDocumentsPermissions + needs to reach the view at all. Model-level only: says nothing + about which documents (if any) this user can actually see. + """ + regular_user.user_permissions.add( + *Permission.objects.filter(codename="view_document"), + ) + return user_client + + @pytest.mark.parametrize( + ("client_fixture", "expected_unrestricted"), + [ + pytest.param("admin_client", True, id="superuser_is_unrestricted"), + pytest.param("viewer_client", False, id="regular_user_is_restricted"), + ], + ) + def test_unrestricted_only_for_superuser( + self, + request: pytest.FixtureRequest, + mocked_stream_chat: mock.MagicMock, + client_fixture: str, + *, + expected_unrestricted: bool, + ) -> None: + """ + GIVEN: + - A superuser, or a regular user holding the global + view_document permission (but no object-level document access) + WHEN: + - They post a chat question with no document_id + THEN: + - stream_chat_with_documents is called with unrestricted=True + only for the superuser; the regular user is always + unrestricted=False, regardless of their permissions + """ + client: APIClient = request.getfixturevalue(client_fixture) + + client.post( + "/api/documents/chat/", + data={"q": "What's in these documents?"}, + format="json", + ) + + assert ( + mocked_stream_chat.call_args.kwargs["unrestricted"] is expected_unrestricted + ) diff --git a/src/documents/views.py b/src/documents/views.py index 4e33d9f9d..4d0a16c4a 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -180,6 +180,7 @@ from documents.permissions import has_system_status_permission from documents.permissions import permitted_document_ids from documents.permissions import permitted_object_ids from documents.permissions import set_permissions_for_object +from documents.permissions import user_is_unrestricted from documents.plugins.date_parsing import get_date_parser from documents.schema import generate_object_with_permissions_schema from documents.search import SearchHit @@ -2329,10 +2330,12 @@ class ChatStreamingView(GenericAPIView[Any]): return HttpResponseForbidden("Insufficient permissions") documents = Document.objects.filter(pk=document.pk) + unrestricted = False else: documents = Document.objects.filter( id__in=permitted_document_ids(request.user), ) + unrestricted = user_is_unrestricted(request.user) output_language = get_llm_output_language( ai_config=ai_config, @@ -2343,6 +2346,7 @@ class ChatStreamingView(GenericAPIView[Any]): stream_chat_with_documents( query_str=question, documents=documents, + unrestricted=unrestricted, output_language=output_language, ), content_type="text/event-stream", diff --git a/src/paperless_ai/chat.py b/src/paperless_ai/chat.py index 03e64af36..a799d4b5c 100644 --- a/src/paperless_ai/chat.py +++ b/src/paperless_ai/chat.py @@ -8,7 +8,8 @@ from documents.models import Document from paperless.config import AIConfig from paperless_ai.client import AIClient from paperless_ai.db import db_connection_released -from paperless_ai.indexing import _document_id_filters +from paperless_ai.indexing import document_id_filters +from paperless_ai.indexing import exclude_document_ids_filter from paperless_ai.indexing import get_rag_prompt_helper from paperless_ai.indexing import load_or_build_index from paperless_ai.indexing import read_store @@ -95,12 +96,15 @@ def _format_chat_metadata_trailer(references: list[dict[str, int | str]]) -> str def stream_chat_with_documents( query_str: str, documents: QuerySet[Document], + *, + unrestricted: bool = False, output_language: str | None = None, ): try: yield from _stream_chat_with_documents( query_str, documents, + unrestricted=unrestricted, output_language=output_language, ) except Exception as e: @@ -111,6 +115,8 @@ def stream_chat_with_documents( def _stream_chat_with_documents( query_str: str, documents: QuerySet[Document], + *, + unrestricted: bool = False, output_language: str | None = None, ): if not documents.exists(): @@ -123,9 +129,18 @@ def _stream_chat_with_documents( from llama_index.core.retrievers import VectorIndexRetriever config = AIConfig() - filters = _document_id_filters( - str(pk) for pk in documents.values_list("pk", flat=True) - ) + 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. + 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) + ) # Hold the shared read lock for the whole operation: the query engine # retrieves from the vector store again during synthesis, so the connection diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index b36a71db7..af20669e5 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -362,7 +362,7 @@ def _embed_nodes(nodes: list["BaseNode"], embed_model) -> None: node.embedding = emb -def _document_id_filters(doc_ids): +def document_id_filters(doc_ids): """Return a MetadataFilters IN filter scoped to ``doc_ids``.""" from llama_index.core.vector_stores.types import FilterOperator from llama_index.core.vector_stores.types import MetadataFilter @@ -396,6 +396,23 @@ def _exclude_document_id_filter(document_id: int | str): ) +def exclude_document_ids_filter(doc_ids): + """Return a MetadataFilters NIN filter excluding every id in ``doc_ids``.""" + from llama_index.core.vector_stores.types import FilterOperator + from llama_index.core.vector_stores.types import MetadataFilter + from llama_index.core.vector_stores.types import MetadataFilters + + return MetadataFilters( + filters=[ + MetadataFilter( + key="document_id", + operator=FilterOperator.NIN, + value=list(doc_ids), + ), + ], + ) + + def update_llm_index( *, iter_wrapper: IterWrapper[Document] = identity, @@ -660,7 +677,7 @@ def retrieve_similar_nodes( filter_parts = [] if allowed_document_ids is not None: - filter_parts.extend(_document_id_filters(allowed_document_ids).filters) + filter_parts.extend(document_id_filters(allowed_document_ids).filters) if document.pk is not None: filter_parts.extend(_exclude_document_id_filter(document.pk).filters) diff --git a/src/paperless_ai/tests/test_chat.py b/src/paperless_ai/tests/test_chat.py index 15c5a545d..be28d0a14 100644 --- a/src/paperless_ai/tests/test_chat.py +++ b/src/paperless_ai/tests/test_chat.py @@ -1,9 +1,14 @@ +from __future__ import annotations + import json +from typing import TYPE_CHECKING +from typing import Any from unittest.mock import MagicMock from unittest.mock import patch import pytest from django.db.models.signals import post_init +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 @@ -18,6 +23,11 @@ from paperless_ai.chat import _build_chat_prompt from paperless_ai.chat import _build_refine_prompt from paperless_ai.chat import stream_chat_with_documents +if TYPE_CHECKING: + from pathlib import Path + + import pytest_mock + @pytest.fixture(autouse=True) def patch_embed_model(): @@ -310,8 +320,40 @@ def test_stream_chat_unexpected_failure_returns_generic_error(caplog) -> None: assert "private provider detail" in caplog.text +def _retriever_filter_values(captured_filters: list[Any]) -> list[str]: + """The value list of the single MetadataFilter the retriever received.""" + assert captured_filters, "VectorIndexRetriever was never constructed" + filt = captured_filters[0] + assert filt is not None, "Retriever must receive a MetadataFilters" + return filt.filters[0].value + + @pytest.mark.django_db class TestStreamChatRetrieval: + @pytest.fixture + def captured_filters(self, mocker: pytest_mock.MockerFixture) -> list[Any]: + """Stub out the AI client and the retriever, capturing the ``filters`` + kwarg of every VectorIndexRetriever construction. + + VectorIndexRetriever is imported inside _stream_chat_with_documents, + so it is patched at the llama_index source for the lazy import to + pick it up. + """ + captured: list[Any] = [] + retriever = mocker.MagicMock() + retriever.retrieve.return_value = [] + + def capture_retriever(*args, **kwargs) -> pytest_mock.MockType: + captured.append(kwargs.get("filters")) + return retriever + + mocker.patch("paperless_ai.chat.AIClient") + mocker.patch( + "llama_index.core.retrievers.VectorIndexRetriever", + side_effect=capture_retriever, + ) + return captured + def test_no_nodes_yields_no_content_message( self, temp_llm_index_dir, @@ -329,9 +371,9 @@ class TestStreamChatRetrieval: def test_chat_filter_contains_only_requested_document_ids( self, - temp_llm_index_dir, - mock_embed_model, - mocker, + temp_llm_index_dir: Path, + mock_embed_model: pytest_mock.MockType, + captured_filters: list[Any], ) -> None: """The MetadataFilter passed to the retriever must be scoped to the requested documents only — content from other indexed documents must @@ -342,22 +384,6 @@ class TestStreamChatRetrieval: indexing.llm_index_add_or_update_document(included) indexing.llm_index_add_or_update_document(excluded) - # VectorIndexRetriever is imported inside _stream_chat_with_documents; - # patch it at the llama_index source so the lazy import picks it up. - captured_filters = [] - mock_retriever = mocker.MagicMock() - mock_retriever.retrieve.return_value = [] - - def capture_retriever(*args, **kwargs): - captured_filters.append(kwargs.get("filters")) - return mock_retriever - - mocker.patch("paperless_ai.chat.AIClient") - mocker.patch( - "llama_index.core.retrievers.VectorIndexRetriever", - side_effect=capture_retriever, - ) - list( chat.stream_chat_with_documents( "question?", @@ -365,13 +391,78 @@ class TestStreamChatRetrieval: ), ) - assert captured_filters, "VectorIndexRetriever was never constructed" - filt = captured_filters[0] - assert filt is not None, "Retriever must receive a MetadataFilters" - filter_values = filt.filters[0].value + filter_values = _retriever_filter_values(captured_filters) assert str(included.pk) in filter_values assert str(excluded.pk) not in filter_values + def test_unrestricted_chat_excludes_nothing_when_no_documents_are_trashed( + self, + temp_llm_index_dir: Path, + mock_embed_model: pytest_mock.MockType, + captured_filters: list[Any], + ) -> None: + """ + GIVEN: + - A document indexed in the vector store, nothing trashed + 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 + """ + document = DocumentFactory.create(content="indexed document content") + indexing.llm_index_add_or_update_document(document) + + list( + chat.stream_chat_with_documents( + "question?", + Document.objects.filter(pk=document.pk), + unrestricted=True, + ), + ) + + assert _retriever_filter_values(captured_filters) == [] + + def test_unrestricted_chat_excludes_trashed_documents( + self, + temp_llm_index_dir: Path, + mock_embed_model: pytest_mock.MockType, + captured_filters: list[Any], + ) -> None: + """ + GIVEN: + - Two indexed documents, one of them trashed -- trashed documents + stay in the vector index until permanently deleted, since + delete_document_from_llm_index is wired to post_delete + WHEN: + - stream_chat_with_documents is called with unrestricted=True + THEN: + - The retriever receives a NOT IN filter excluding the trashed + document's id, so an unrestricted caller (e.g. a superuser) + never has trashed content surfaced in a chat answer + """ + kept = DocumentFactory.create(content="kept document content") + trashed = DocumentFactory.create(content="trashed document content") + indexing.llm_index_add_or_update_document(kept) + indexing.llm_index_add_or_update_document(trashed) + Document.global_objects.filter(pk=trashed.pk).update( + deleted_at=timezone.now(), + ) + + list( + chat.stream_chat_with_documents( + "question?", + Document.objects.filter(pk=kept.pk), + unrestricted=True, + ), + ) + + filter_values = _retriever_filter_values(captured_filters) + assert str(trashed.pk) in filter_values + assert str(kept.pk) not in filter_values + @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 3ed1b7e42..1c4dbf2f9 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -1,5 +1,6 @@ import inspect import sqlite3 +from collections.abc import Callable from collections.abc import Generator from pathlib import Path @@ -97,6 +98,18 @@ def _ne_filter(document_id: int): ) +def _nin_filter(document_ids: list[int]): + return MetadataFilters( + filters=[ + MetadataFilter( + key="document_id", + operator=FilterOperator.NIN, + value=document_ids, + ), + ], + ) + + class TestCrud: def test_add_then_query_returns_node(self, store) -> None: node = make_node("n1", 1) @@ -280,6 +293,47 @@ class TestBuildWhere: "b1", ] + def test_nin_filter_translates_to_not_in_clause(self) -> None: + where, params = _build_where(_nin_filter([1, 2])) + assert where == "(document_id NOT IN (?,?))" + assert params == [1, 2] + + def test_query_with_nin_filter_excludes_matching_documents(self, store) -> None: + store.add([make_node("a1", 1), make_node("b1", 2), make_node("c1", 3)]) + assert sorted( + _query(store, [0.0] * DIM, top_k=5, filters=_nin_filter([1, 2])).ids, + ) == ["c1"] + + def test_empty_in_filter_excludes_everything(self) -> None: + """ + GIVEN: + - An IN filter with an empty value list + WHEN: + - _build_where() translates it to SQL + THEN: + - It excludes everything (the opposite of an empty NOT IN + filter) -- an empty inclusion list must never widen results + """ + where, params = _build_where(_in_filter([])) + assert where == "(1 = 0)" + assert params == [] + + def test_empty_nin_filter_excludes_nothing(self) -> None: + """ + GIVEN: + - A NOT IN filter with an empty value list -- e.g. an + unrestricted chat caller when nothing is currently trashed + WHEN: + - _build_where() translates it to SQL + THEN: + - It excludes nothing (unlike an empty IN filter, which + excludes everything) -- an empty exclusion list must never + narrow results + """ + where, params = _build_where(_nin_filter([])) + assert where == "(1 = 1)" + assert params == [] + def test_fails_closed_when_no_filter_is_translatable(self) -> None: # A nested MetadataFilters is not a MetadataFilter, so it is skipped. # With no translatable clauses, the function must fail closed rather @@ -297,24 +351,31 @@ class TestBuildWhere: assert where == "1 = 0" assert params == [] - def test_fails_closed_when_in_filter_exceeds_max_values( + @pytest.mark.parametrize( + "build_filter", + [_in_filter, _nin_filter], + ids=["in", "nin"], + ) + def test_fails_closed_when_filter_exceeds_max_values( self, + build_filter: Callable[[list[str]], MetadataFilters], caplog: pytest.LogCaptureFixture, ) -> None: """ GIVEN: - - An 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) + - 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) WHEN: - _build_where() translates it to SQL THEN: - - It fails closed ("1 = 0", no params) instead of building an - IN 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 + - 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 """ - oversized = _in_filter([str(i) for i in range(_MAX_IN_VALUES + 1)]) + oversized = build_filter([str(i) for i in range(_MAX_IN_VALUES + 1)]) with caplog.at_level("WARNING"): where, params = _build_where(oversized) diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index b7de0936d..c7675261d 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -107,12 +107,13 @@ def _vec0_params(rows: list[_Row]) -> list[tuple[str, int, str, bytes]]: def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: - """Translate the EQ / IN / 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 INTEGER; - every value is coerced via int() here so callers (which today still pass - strings in places, e.g. indexing.py's MetadataFilter construction) don't - have to be individually correct -- vec0 doesn't coerce types itself. + """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 + INTEGER; every value is coerced via int() here so callers (which today + still pass strings in places, e.g. indexing.py's MetadataFilter + construction) don't have to be individually correct -- vec0 doesn't + coerce types itself. """ if filters is None or not filters.filters: return "", [] @@ -125,20 +126,25 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: continue if f.key not in _FILTER_COLUMNS: # pragma: no cover - we build the keys raise NotImplementedError(f"Unsupported filter column: {f.key}") - if f.operator == FilterOperator.IN: + if f.operator in (FilterOperator.IN, FilterOperator.NIN): + is_in = f.operator == FilterOperator.IN + sql_op = "IN" if is_in else "NOT IN" values = [int(v) for v in f.value] # type: ignore[union-attr] - if not values: # pragma: no cover - clauses.append("1 = 0") + if not values: + # An empty IN list matches nothing; an empty NOT IN list + # excludes nothing, so it matches everything. + clauses.append("1 = 0" if is_in else "1 = 1") continue if len(values) > _MAX_IN_VALUES: - # Fail closed (see the empty-clauses case below) rather than - # let SQLite raise "too many SQL variables" past its own - # limit: this filter scopes document access, so an IN list - # too large to safely bind must match no rows, never widen - # the scope to "everything" by accident. + # 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 an IN filter on %r with %d values " + "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, @@ -146,7 +152,7 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: clauses.append("1 = 0") continue placeholders = ",".join("?" for _ in values) - clauses.append(f"{f.key} IN ({placeholders})") + clauses.append(f"{f.key} {sql_op} ({placeholders})") params.extend(values) elif f.operator == FilterOperator.EQ: clauses.append(f"{f.key} = ?") @@ -154,7 +160,7 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: elif f.operator == FilterOperator.NE: clauses.append(f"{f.key} != ?") params.append(int(f.value)) - else: # pragma: no cover - we only ever build EQ/IN/NE filters + else: # pragma: no cover - we only ever build EQ/IN/NIN/NE filters raise NotImplementedError(f"Unsupported filter operator: {f.operator}") if not clauses: # Filters were requested but none could be translated. Fail closed