From e4e48733e100a4ae11c130e30ff176e6b88b9b13 Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Thu, 13 Aug 2026 13:41:30 -0700 Subject: [PATCH] - get_assigned_metadata() now permission-filters the document's own tags/type/correspondent/storage_path per object - existing_ids the model returns are now restricted to ids that were actually offered as candidates in the prompt - ai_suggestions no longer skips permission filtering on a cache hit. The cache now stores the raw existing_ids/new_names choices rather than resolved object ids - ai_suggestions deduplicates matched objects by id - Performance updates for keeping things into a queryset instead --- src/documents/permissions.py | 60 ++++++--- .../test_permission_filtering_security.py | 60 ++++++--- src/documents/tests/test_views.py | 118 +++++++++++++++++- src/documents/views.py | 85 ++++++++----- src/paperless_ai/ai_classifier.py | 57 ++++++++- src/paperless_ai/base_model.py | 2 +- src/paperless_ai/matching.py | 11 +- src/paperless_ai/taxonomy.py | 70 ++++++++--- src/paperless_ai/tests/test_ai_classifier.py | 105 +++++++++++++++- src/paperless_ai/tests/test_base_model.py | 2 +- src/paperless_ai/tests/test_taxonomy.py | 75 ++++++++++- 11 files changed, 540 insertions(+), 105 deletions(-) diff --git a/src/documents/permissions.py b/src/documents/permissions.py index 82584d846..d22882753 100644 --- a/src/documents/permissions.py +++ b/src/documents/permissions.py @@ -1,4 +1,5 @@ from typing import Any +from typing import TypeVar from django.contrib.auth.models import Group from django.contrib.auth.models import Permission @@ -235,35 +236,56 @@ def permitted_object_ids( ).values_list("id", flat=True) -def visible_object_ids_or_none( - user: User | None, - model: type[Model], - perm: str, -) -> set[int] | None: +ModelT = TypeVar("ModelT", bound=Model) + + +def user_is_unrestricted(user: User | None) -> bool: """ - Return the set of object IDs of ``model`` that ``user`` may see with - ``perm``, or ``None`` meaning "no restriction at all". + True when ``user`` means "no restriction at all" (an absent user, or an + *active* superuser) without needing a database check to know it. - ``None`` is returned only for an absent user or an *active* superuser. ``permitted_object_ids(None, ...)`` itself means the much narrower "only unowned rows", which is NOT the same thing as "no user filtering - requested", so that case has to be special-cased before ever calling it. + requested", so callers must special-case this before ever calling it. + A deactivated superuser is deliberately NOT unrestricted here, matching + permitted_object_ids's own is_active-before-is_superuser ordering. - Every other case is delegated to ``permitted_object_ids`` rather than - re-deciding here, so its ordering is inherited instead of duplicated: a - deactivated superuser must NOT be handed "no restriction", it gets an - empty set (nothing visible), and an unauthenticated user still gets the - unowned rows. + Callers that can avoid a database round trip entirely when this is true + (e.g. checking a single already-loaded object's visibility rather than + filtering a queryset) should do so via this function directly, rather + than through restrict_queryset_to_visible() below. """ if user is None: - return None - if ( + return True + return ( getattr(user, "is_authenticated", False) and getattr(user, "is_active", False) and getattr(user, "is_superuser", False) - ): - return None - return set(permitted_object_ids(user, model, perm)) + ) + + +def restrict_queryset_to_visible( + queryset: QuerySet[ModelT], + user: User | None, + perm: str, +) -> QuerySet[ModelT]: + """ + Restrict ``queryset`` to the rows ``user`` may see with ``perm``. + + Delegates the visibility check to the database as a + ``WHERE id IN (subquery)`` rather than materializing the full + permitted-id set into a Python collection first: a caller that only + needs to check a small handful of rows (a resolved-id list, a few + RAG-neighbour candidate ids) never pays for scanning or holding the + installation's entire taxonomy in memory to do it. + + Returns ``queryset`` unchanged for user_is_unrestricted(user); every + other case is delegated to ``permitted_object_ids`` rather than + re-deciding the ordering here. + """ + if user_is_unrestricted(user): + return queryset + return queryset.filter(pk__in=permitted_object_ids(user, queryset.model, perm)) def permitted_document_ids( diff --git a/src/documents/tests/test_permission_filtering_security.py b/src/documents/tests/test_permission_filtering_security.py index c00b75f6c..282e25420 100644 --- a/src/documents/tests/test_permission_filtering_security.py +++ b/src/documents/tests/test_permission_filtering_security.py @@ -22,7 +22,7 @@ from documents.models import StoragePath from documents.models import Tag from documents.permissions import permitted_document_ids from documents.permissions import permitted_object_ids -from documents.permissions import visible_object_ids_or_none +from documents.permissions import restrict_queryset_to_visible from documents.serialisers import _get_viewable_duplicates from documents.tests.factories import CorrespondentFactory from documents.tests.factories import DocumentFactory @@ -737,7 +737,7 @@ class TestBulkEditObjectsTagDescendantPartialPermission: NOTE: this uses ``set_permissions`` (owner reassignment) rather than ``delete`` as the operation, because Tag.tn_parent (django-treenode) cascades deletes to descendants at the database/ORM level regardless - of which tags the view resolved into ``objs`` -- a delete-based test + of which tags the view resolved into ``objs`` - a delete-based test would pass/fail based on FK cascade behavior, not on whether the descendant-expansion logic itself respected per-object permissions. """ @@ -787,46 +787,58 @@ class TestBulkEditObjectsTagDescendantPartialPermission: @pytest.mark.django_db -class TestVisibleObjectIdsOrNone: - """``None`` from visible_object_ids_or_none() means "no restriction at - all", so the cases that may return it have to be kept narrow.""" +class TestRestrictQuerysetToVisible: + """restrict_queryset_to_visible() returns its queryset argument + unchanged only for "no restriction at all", so the cases that may do + that have to be kept narrow.""" def test_no_user_means_no_restriction(self) -> None: """ GIVEN: - No user at all (a system-triggered call) WHEN: - - visible_object_ids_or_none() is called + - restrict_queryset_to_visible() is called THEN: - - None is returned, i.e. no filtering, rather than + - The queryset is returned unfiltered, rather than permitted_object_ids(None, ...)'s narrower "unowned rows only" """ owner = User.objects.create_user(username="vis_none_owner") - TagFactory(owner=owner) + tag = TagFactory(owner=owner) - assert visible_object_ids_or_none(None, Tag, "view_tag") is None + visible = restrict_queryset_to_visible(Tag.objects.all(), None, "view_tag") + + assert tag.pk in visible.values_list("pk", flat=True) def test_active_superuser_means_no_restriction(self) -> None: """ GIVEN: - An active superuser WHEN: - - visible_object_ids_or_none() is called + - restrict_queryset_to_visible() is called THEN: - - None is returned, skipping the permission lookup entirely + - The queryset is returned unfiltered, skipping the permission + lookup entirely """ superuser = User.objects.create_superuser(username="vis_active_super") + owner = User.objects.create_user(username="vis_active_super_owner") + tag = TagFactory(owner=owner) - assert visible_object_ids_or_none(superuser, Tag, "view_tag") is None + visible = restrict_queryset_to_visible( + Tag.objects.all(), + superuser, + "view_tag", + ) + + assert tag.pk in visible.values_list("pk", flat=True) def test_inactive_superuser_is_denied_not_unrestricted(self) -> None: """ GIVEN: - A deactivated superuser WHEN: - - visible_object_ids_or_none() is called + - restrict_queryset_to_visible() is called THEN: - - An empty set (nothing visible) is returned, never None -- + - No rows are visible, never the whole unrestricted queryset - deactivation has to win over the superuser shortcut, matching permitted_object_ids's own ordering """ @@ -838,23 +850,31 @@ class TestVisibleObjectIdsOrNone: TagFactory(owner=None) TagFactory(owner=user) - assert visible_object_ids_or_none(user, Tag, "view_tag") == set() + visible = restrict_queryset_to_visible(Tag.objects.all(), user, "view_tag") + + assert not visible.exists() def test_regular_user_gets_permitted_ids(self) -> None: """ GIVEN: - An ordinary active user and a tag owned by someone else WHEN: - - visible_object_ids_or_none() is called + - restrict_queryset_to_visible() is called THEN: - - Only the ids permitted_object_ids() reports are returned + - Only the rows permitted_object_ids() reports are visible """ user = User.objects.create_user(username="vis_regular") other = User.objects.create_user(username="vis_regular_other") own = TagFactory(owner=user) hidden = TagFactory(owner=other) - visible = visible_object_ids_or_none(user, Tag, "view_tag") + visible_ids = set( + restrict_queryset_to_visible( + Tag.objects.all(), + user, + "view_tag", + ).values_list("pk", flat=True), + ) - assert own.pk in visible - assert hidden.pk not in visible + assert own.pk in visible_ids + assert hidden.pk not in visible_ids diff --git a/src/documents/tests/test_views.py b/src/documents/tests/test_views.py index 650c96c35..81331f1ff 100644 --- a/src/documents/tests/test_views.py +++ b/src/documents/tests/test_views.py @@ -352,20 +352,95 @@ class TestAISuggestions(DirectoriesMixin, TestCase): mock_refresh_cache, mock_get_cache, ) -> None: - mock_get_cache.return_value = MagicMock(suggestions={"tags": ["tag1", "tag2"]}) + """ + GIVEN: + - A cached LLM classification holding the raw existing_ids/ + new_names choices (never resolved object ids) + WHEN: + - ai_suggestions is requested + THEN: + - The cached choices are resolved into ids for this request + (not returned verbatim from the cache) and the cache's TTL is + refreshed + """ + mock_get_cache.return_value = MagicMock( + suggestions={ + "title": "Cached Title", + "tags": {"existing_ids": [self.tag1.pk], "new_names": []}, + "correspondents": {"existing_ids": [], "new_names": []}, + "document_types": {"existing_ids": [], "new_names": []}, + "storage_paths": {"existing_ids": [], "new_names": []}, + "dates": [], + }, + ) self.client.force_login(user=self.user) response = self.client.get( f"/api/documents/{self.document.pk}/ai_suggestions/", ) self.assertEqual(response.status_code, status.HTTP_200_OK) - self.assertEqual(response.json(), {"tags": ["tag1", "tag2"]}) + self.assertEqual(response.json()["title"], "Cached Title") + self.assertEqual(response.json()["tags"], [self.tag1.pk]) mock_get_cache.assert_called_once_with( self.document.pk, backend="mock_backend", ) mock_refresh_cache.assert_called_once_with(self.document.pk) + @patch("documents.views.get_llm_suggestion_cache") + @patch("documents.views.refresh_suggestions_cache") + @override_settings( + AI_ENABLED=True, + LLM_BACKEND="mock_backend", + ) + def test_ai_suggestions_cache_hit_re_filters_for_narrower_requester( + self, + mock_refresh_cache, + mock_get_cache, + ) -> None: + """ + GIVEN: + - A cached LLM classification whose existing_ids include a tag + only visible to a broader-visibility user (e.g. the requester + who originally generated it) + - A second, non-superuser requester who may change the document + but has no permission to view that tag + WHEN: + - ai_suggestions is requested by the second requester and the + cache is hit + THEN: + - The cache hit still runs permission filtering fresh for this + requester; the invisible tag id does not leak into either the + matched or suggested tags + """ + tag_owner = User.objects.create_user(username="cache_tag_owner") + invisible_tag = Tag.objects.create(name="cache_restricted", owner=tag_owner) + requester = User.objects.create_user(username="cache_requester") + requester.user_permissions.add( + *Permission.objects.filter( + codename__in=["view_document", "change_document", "view_tag"], + ), + ) + mock_get_cache.return_value = MagicMock( + suggestions={ + "title": "Untitled", + "tags": {"existing_ids": [invisible_tag.pk], "new_names": []}, + "correspondents": {"existing_ids": [], "new_names": []}, + "document_types": {"existing_ids": [], "new_names": []}, + "storage_paths": {"existing_ids": [], "new_names": []}, + "dates": [], + }, + ) + + self.client.force_login(user=requester) + response = self.client.get( + f"/api/documents/{self.document.pk}/ai_suggestions/", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.json()["tags"], []) + self.assertEqual(response.json()["suggested_tags"], []) + @patch("documents.views.get_ai_document_classification") @override_settings( AI_ENABLED=True, @@ -623,6 +698,45 @@ class TestAISuggestions(DirectoriesMixin, TestCase): self.assertEqual(response.json()["tags"], [self.tag1.pk]) self.assertEqual(response.json()["suggested_tags"], ["Follow-up"]) + @patch("documents.views.get_ai_document_classification") + @override_settings( + AI_ENABLED=True, + LLM_BACKEND="mock_backend", + ) + def test_ai_suggestions_deduplicates_id_matched_via_both_paths( + self, + mock_get_ai_classification, + ) -> None: + """ + GIVEN: + - AI classification returns the same tag both as an existing_id + and as a new_name that fuzzy-matches that same tag + WHEN: + - ai_suggestions is requested + THEN: + - The tag's id appears exactly once in the response, not twice + """ + mock_get_ai_classification.return_value = { + "title": "Lab Report", + "tags": { + "existing_ids": [self.tag1.pk], + "new_names": [self.tag1.name], + }, + "correspondents": {"existing_ids": [], "new_names": []}, + "document_types": {"existing_ids": [], "new_names": []}, + "storage_paths": {"existing_ids": [], "new_names": []}, + "dates": [], + } + + self.client.force_login(user=self.user) + response = self.client.get( + f"/api/documents/{self.document.pk}/ai_suggestions/", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.json()["tags"], [self.tag1.pk]) + self.assertEqual(response.json()["suggested_tags"], []) + @patch("documents.views.get_ai_document_classification") @override_settings( AI_ENABLED=True, diff --git a/src/documents/views.py b/src/documents/views.py index f4edccbcc..68f786b6b 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -1554,34 +1554,48 @@ class DocumentViewSet( ) if cached_llm_suggestions: + # Only the raw model choices are cached, never resolved object + # ids. resolve_choice() below still runs permission filtering + # freshly for this requester on every request, cache hit or not, + # so a resolved id cached for one user's visibility can never be + # handed unfiltered to a second, less-privileged requester of + # the same (backend-keyed, not user-keyed) cache entry. refresh_suggestions_cache(doc.pk) - return Response(cached_llm_suggestions.suggestions) - - try: - llm_suggestions = get_ai_document_classification( - doc, - request.user, - output_language, - ) - except ValueError as exc: - logger.exception( - "Invalid AI configuration while generating suggestions for " - "document %s: %s", + llm_suggestions = cached_llm_suggestions.suggestions + else: + try: + llm_suggestions = get_ai_document_classification( + doc, + request.user, + output_language, + ) + except ValueError as exc: + logger.exception( + "Invalid AI configuration while generating suggestions for " + "document %s: %s", + doc.pk, + exc, + exc_info=True, + ) + raise ValidationError( + {"ai": [_("Invalid AI configuration.")]}, + ) from exc + except LLMTimeoutError as exc: + logger.exception( + "AI backend timed out while generating suggestions for " + "document %s: %s", + doc.pk, + exc, + exc_info=True, + ) + return Response( + {"ai": [_("AI backend request timed out.")]}, + status=status.HTTP_503_SERVICE_UNAVAILABLE, + ) + set_llm_suggestions_cache( doc.pk, - exc, - exc_info=True, - ) - raise ValidationError({"ai": [_("Invalid AI configuration.")]}) from exc - except LLMTimeoutError as exc: - logger.exception( - "AI backend timed out while generating suggestions for document %s: %s", - doc.pk, - exc, - exc_info=True, - ) - return Response( - {"ai": [_("AI backend request timed out.")]}, - status=status.HTTP_503_SERVICE_UNAVAILABLE, + llm_suggestions, + backend=llm_cache_backend, ) tags_choice: TaxonomyChoiceDict = llm_suggestions["tags"] @@ -1595,11 +1609,24 @@ class DocumentViewSet( match_names: Callable[[list[str], User], list], ) -> list: """The ids the model picked from the candidates it was shown, plus - name matches for the values it proposed as new.""" - return resolve_ids(choice["existing_ids"], request.user) + match_names( + name matches for the values it proposed as new. The schema allows + the same object to satisfy both an existing_id and a new_name in + one valid response, so results are deduplicated by pk (keeping + first-seen order) rather than trusting the two lookups to be + disjoint. + """ + matched = resolve_ids(choice["existing_ids"], request.user) + match_names( choice["new_names"], request.user, ) + seen_ids: set[int] = set() + deduped = [] + for obj in matched: + if obj.pk in seen_ids: + continue + seen_ids.add(obj.pk) + deduped.append(obj) + return deduped matched_tags = resolve_choice( tags_choice, @@ -1647,8 +1674,6 @@ class DocumentViewSet( "dates": llm_suggestions["dates"], } - set_llm_suggestions_cache(doc.pk, resp_data, backend=llm_cache_backend) - return Response(resp_data) @action(methods=["get"], detail=True, filter_backends=[]) diff --git a/src/paperless_ai/ai_classifier.py b/src/paperless_ai/ai_classifier.py index d27f0c207..952e09cf3 100644 --- a/src/paperless_ai/ai_classifier.py +++ b/src/paperless_ai/ai_classifier.py @@ -159,7 +159,7 @@ def get_taxonomy_context( propagating the exception - a vector-store outage should not block classification, only its RAG-assisted enrichment. """ - assigned = get_assigned_metadata(document) + assigned = get_assigned_metadata(document, user) try: visible_document_ids = ( None @@ -220,6 +220,49 @@ def parse_ai_response(raw: dict) -> ClassificationSuggestions: ) +def _restrict_to_shown_candidates( + suggestions: ClassificationSuggestions, + candidates: TaxonomyCandidates, +) -> ClassificationSuggestions: + """Drop any existing_id the model returned that was never actually + offered as a candidate in the prompt. The response schema permits any + integer, so a hallucinated id could otherwise silently resolve to a + real, visible, but completely unrelated object - this keeps + "reused an existing value" a fact about what the model was actually + shown, not just about what integer it happened to emit. When no + candidates were shown in a category at all (or the field was omitted + from the response), every existing_id in that category is dropped; + new_names is never touched here. + """ + + def _restrict(choice: TaxonomyChoiceDict, shown: set[int]) -> TaxonomyChoiceDict: + return TaxonomyChoiceDict( + existing_ids=[i for i in choice["existing_ids"] if i in shown], + new_names=choice["new_names"], + ) + + return ClassificationSuggestions( + title=suggestions["title"], + tags=_restrict( + suggestions["tags"], + {c["id"] for c in candidates["tags"]}, + ), + correspondents=_restrict( + suggestions["correspondents"], + {c["id"] for c in candidates["correspondents"]}, + ), + document_types=_restrict( + suggestions["document_types"], + {c["id"] for c in candidates["document_types"]}, + ), + storage_paths=_restrict( + suggestions["storage_paths"], + {c["id"] for c in candidates["storage_paths"]}, + ), + dates=suggestions["dates"], + ) + + def get_ai_document_classification( document: Document, user: User | None = None, @@ -237,11 +280,12 @@ def get_ai_document_classification( context=context, ) else: + candidates = empty_taxonomy_candidates() prompt = build_prompt_without_rag( document, ai_config, - candidates=empty_taxonomy_candidates(), - assigned=get_assigned_metadata(document), + candidates=candidates, + assigned=get_assigned_metadata(document, user), ) client = AIClient() @@ -249,7 +293,10 @@ def get_ai_document_classification( # is not pinned for the call's duration; see paperless_ai.db and #12976. with db_connection_released(): result = client.run_llm_query(prompt) - suggestions = parse_ai_response(result) + suggestions = _restrict_to_shown_candidates( + parse_ai_response(result), + candidates, + ) if output_language: localized = client.run_llm_query( build_localization_prompt(suggestions, output_language), @@ -257,7 +304,7 @@ def get_ai_document_classification( localized_suggestions = parse_ai_response(localized) def _localized_choice(field: str) -> TaxonomyChoiceDict: - # existing_ids always come from the ORIGINAL suggestions -- + # existing_ids always come from the ORIGINAL suggestions - # never from localized_suggestions, whatever the model echoed # back there. This is the concrete fix for the bug this # feature exists to close: localization must never be able to diff --git a/src/paperless_ai/base_model.py b/src/paperless_ai/base_model.py index 8484527a0..3df311035 100644 --- a/src/paperless_ai/base_model.py +++ b/src/paperless_ai/base_model.py @@ -39,7 +39,7 @@ class TaxonomyChoiceDict(TypedDict): class ClassificationSuggestions(TypedDict): - """Plain-dict counterpart of DocumentClassifierSchema.model_dump() -- + """Plain-dict counterpart of DocumentClassifierSchema.model_dump() - the shape threaded through parse_ai_response, build_localization_prompt, get_ai_document_classification, and the ai_suggestions view.""" diff --git a/src/paperless_ai/matching.py b/src/paperless_ai/matching.py index 30feae106..0cadaf36c 100644 --- a/src/paperless_ai/matching.py +++ b/src/paperless_ai/matching.py @@ -12,7 +12,7 @@ from documents.models import DocumentType from documents.models import StoragePath from documents.models import Tag from documents.permissions import get_objects_for_user_owner_aware -from documents.permissions import visible_object_ids_or_none +from documents.permissions import restrict_queryset_to_visible MATCH_THRESHOLD = 0.8 @@ -34,10 +34,11 @@ def _resolve_visible_ids( """ if not ids: return [] - visible_ids = visible_object_ids_or_none(user, model, perm) - queryset = model.objects.filter(pk__in=ids) - if visible_ids is not None: - queryset = queryset.filter(pk__in=visible_ids) + queryset = restrict_queryset_to_visible( + model.objects.filter(pk__in=ids), + user, + perm, + ) return list(queryset) diff --git a/src/paperless_ai/taxonomy.py b/src/paperless_ai/taxonomy.py index 7e639a942..b9412337c 100644 --- a/src/paperless_ai/taxonomy.py +++ b/src/paperless_ai/taxonomy.py @@ -12,7 +12,8 @@ from documents.models import Document from documents.models import DocumentType from documents.models import StoragePath from documents.models import Tag -from documents.permissions import visible_object_ids_or_none +from documents.permissions import restrict_queryset_to_visible +from documents.permissions import user_is_unrestricted if TYPE_CHECKING: from llama_index.core.schema import NodeWithScore @@ -53,17 +54,50 @@ def empty_taxonomy_candidates() -> TaxonomyCandidates: ) -def get_assigned_metadata(document: Document) -> AssignedMetadata: +def _visible_name( + obj: Model | None, + user: User | None, + perm: str, +) -> str | None: + """``obj``'s name if ``user`` may see it under ``perm``, else None - a + document being visible to a user does not imply every object assigned to + it is (per-object guardian permissions can differ), so each assigned + relation is checked individually rather than trusted because it's + already sitting on a document this user can open. + + Checks user_is_unrestricted() before ever touching type(obj).objects, so + the common "no restriction" case (no user, or an active superuser) never + needs obj to be backed by a real queryable row. + """ + if obj is None: + return None + if user_is_unrestricted(user): + return obj.name + visible = restrict_queryset_to_visible( + type(obj).objects.filter(pk=obj.pk), + user, + perm, + ) + return obj.name if visible.exists() else None + + +def get_assigned_metadata(document: Document, user: User | None) -> AssignedMetadata: """The document's own current taxonomy. Authoritative context, not a candidate list - the model is never asked to add, remove, or replace these values, only to use them when helpful for the title and for fields that are still empty. + + Permission-filtered the same way build_taxonomy_candidates() is: a + document a user may change/view does not imply every tag/type/ + correspondent/storage_path assigned to it is visible to that same user, + so names the user cannot see are never surfaced into the prompt. """ + visible_tags = restrict_queryset_to_visible(document.tags.all(), user, "view_tag") return AssignedMetadata( - tags=sorted(tag.name for tag in document.tags.all()), - document_type=document.document_type.name if document.document_type else None, - correspondent=document.correspondent.name if document.correspondent else None, - storage_path=document.storage_path.name if document.storage_path else None, + tags=sorted(tag.name for tag in visible_tags), + document_type=_visible_name(document.document_type, user, "view_documenttype"), + correspondent=_visible_name(document.correspondent, user, "view_correspondent"), + storage_path=_visible_name(document.storage_path, user, "view_storagepath"), ) @@ -91,17 +125,21 @@ def _visible_ranked_candidates( limit: int, ) -> list[TaxonomyCandidate]: """Drop anything ``user`` may not see, resolve the survivors' names, and - return them ranked by descending weight and capped at ``limit``.""" - visible_ids = visible_object_ids_or_none(user, model, perm) - if visible_ids is not None: - weighted_ids = { - object_id: weight - for object_id, weight in weighted_ids.items() - if object_id in visible_ids - } - id_to_name = dict( - model.objects.filter(pk__in=weighted_ids).values_list("id", "name"), + return them ranked by descending weight and capped at ``limit``. + + The visibility check restricts the query to just this small + weighted_ids set rather than materializing every id `user` may see + installation-wide - resolving names and checking visibility is one + query either way, so this never pays for scanning the whole taxonomy. + """ + if not weighted_ids: + return [] + visible_queryset = restrict_queryset_to_visible( + model.objects.filter(pk__in=weighted_ids), + user, + perm, ) + id_to_name = dict(visible_queryset.values_list("id", "name")) candidates = [ TaxonomyCandidate(id=object_id, name=id_to_name[object_id], weight=weight) for object_id, weight in weighted_ids.items() diff --git a/src/paperless_ai/tests/test_ai_classifier.py b/src/paperless_ai/tests/test_ai_classifier.py index 765eda228..fe8c12e88 100644 --- a/src/paperless_ai/tests/test_ai_classifier.py +++ b/src/paperless_ai/tests/test_ai_classifier.py @@ -11,12 +11,18 @@ from documents.tests.factories import DocumentFactory from documents.tests.factories import TagFactory from documents.tests.factories import UserFactory from paperless.config import AIConfig +from paperless_ai.ai_classifier import _restrict_to_shown_candidates from paperless_ai.ai_classifier import build_localization_prompt from paperless_ai.ai_classifier import build_prompt_with_rag from paperless_ai.ai_classifier import build_prompt_without_rag from paperless_ai.ai_classifier import get_ai_document_classification from paperless_ai.ai_classifier import get_language_name from paperless_ai.ai_classifier import get_taxonomy_context +from paperless_ai.base_model import ClassificationSuggestions +from paperless_ai.base_model import TaxonomyChoiceDict +from paperless_ai.taxonomy import TaxonomyCandidate +from paperless_ai.taxonomy import TaxonomyCandidates +from paperless_ai.taxonomy import empty_taxonomy_candidates @pytest.fixture @@ -603,14 +609,22 @@ def test_build_prompt_without_rag_identical_when_no_hints(): @pytest.mark.django_db @patch("paperless_ai.ai_classifier.AIClient") +@patch("paperless_ai.ai_classifier.build_taxonomy_candidates") @patch("paperless_ai.ai_classifier.retrieve_similar_nodes") +@override_settings( + LLM_EMBEDDING_BACKEND="huggingface", + LLM_BACKEND="ollama", + LLM_MODEL="some_model", +) def test_get_ai_document_classification_localizes_only_new_names( mock_retrieve, + mock_build_candidates, mock_client_cls, ): """ GIVEN: - - A classification response with a resolved existing tag id + - A classification response with a resolved existing tag id that + was actually offered as a candidate - A localization response that echoes back a different existing_ids value WHEN: - get_ai_document_classification() is called with an output_language @@ -621,6 +635,12 @@ def test_get_ai_document_classification_localizes_only_new_names( """ document = DocumentFactory.create(content="Some content") mock_retrieve.return_value = [] + mock_build_candidates.return_value = TaxonomyCandidates( + tags=[TaxonomyCandidate(id=12, name="Contractor", weight=1.0)], + document_types=[], + correspondents=[], + storage_paths=[], + ) mock_client = mock_client_cls.return_value mock_client.run_llm_query.side_effect = [ { @@ -649,3 +669,86 @@ def test_get_ai_document_classification_localizes_only_new_names( assert "Contractor Work" in localization_prompt assert result["tags"]["existing_ids"] == [12] # untouched by localization assert result["tags"]["new_names"] == ["Auftragsarbeit"] + + +class TestRestrictToShownCandidates: + def test_hallucinated_id_not_among_candidates_is_dropped(self) -> None: + """ + GIVEN: + - A tag candidate shown to the model with id=12 + - A model response with existing_ids=[12, 999] for tags, where + 999 was never offered as a candidate + WHEN: + - _restrict_to_shown_candidates() is called + THEN: + - Only the id that was actually shown survives; the hallucinated + id is dropped rather than being trusted to resolve to whatever + real, visible, unrelated object it happens to match + """ + suggestions = ClassificationSuggestions( + title="T", + tags=TaxonomyChoiceDict(existing_ids=[12, 999], new_names=[]), + correspondents=TaxonomyChoiceDict(existing_ids=[], new_names=[]), + document_types=TaxonomyChoiceDict(existing_ids=[], new_names=[]), + storage_paths=TaxonomyChoiceDict(existing_ids=[], new_names=[]), + dates=[], + ) + candidates = TaxonomyCandidates( + tags=[TaxonomyCandidate(id=12, name="Contractor", weight=1.0)], + document_types=[], + correspondents=[], + storage_paths=[], + ) + + result = _restrict_to_shown_candidates(suggestions, candidates) + + assert result["tags"]["existing_ids"] == [12] + + def test_no_candidates_shown_drops_every_existing_id(self) -> None: + """ + GIVEN: + - No candidates were shown in any category + - A model response with existing_ids populated anyway + WHEN: + - _restrict_to_shown_candidates() is called + THEN: + - Every existing_id is dropped across all four categories - an + id can only be trusted if the prompt actually offered it + """ + suggestions = ClassificationSuggestions( + title="T", + tags=TaxonomyChoiceDict(existing_ids=[1], new_names=[]), + correspondents=TaxonomyChoiceDict(existing_ids=[2], new_names=[]), + document_types=TaxonomyChoiceDict(existing_ids=[3], new_names=[]), + storage_paths=TaxonomyChoiceDict(existing_ids=[4], new_names=[]), + dates=[], + ) + + result = _restrict_to_shown_candidates(suggestions, empty_taxonomy_candidates()) + + assert result["tags"]["existing_ids"] == [] + assert result["correspondents"]["existing_ids"] == [] + assert result["document_types"]["existing_ids"] == [] + assert result["storage_paths"]["existing_ids"] == [] + + def test_new_names_are_never_touched(self) -> None: + """ + GIVEN: + - A model response with new_names populated + WHEN: + - _restrict_to_shown_candidates() is called + THEN: + - new_names passes through unchanged regardless of candidates + """ + suggestions = ClassificationSuggestions( + title="T", + tags=TaxonomyChoiceDict(existing_ids=[], new_names=["Brand New Tag"]), + correspondents=TaxonomyChoiceDict(existing_ids=[], new_names=[]), + document_types=TaxonomyChoiceDict(existing_ids=[], new_names=[]), + storage_paths=TaxonomyChoiceDict(existing_ids=[], new_names=[]), + dates=[], + ) + + result = _restrict_to_shown_candidates(suggestions, empty_taxonomy_candidates()) + + assert result["tags"]["new_names"] == ["Brand New Tag"] diff --git a/src/paperless_ai/tests/test_base_model.py b/src/paperless_ai/tests/test_base_model.py index dda5d6411..0ffda52cf 100644 --- a/src/paperless_ai/tests/test_base_model.py +++ b/src/paperless_ai/tests/test_base_model.py @@ -50,7 +50,7 @@ def test_document_classifier_schema_json_schema_is_self_contained(): client.py hands this generated schema straight to the LLM backend as the response-format constraint (Ollama's format=json_schema, and the OpenAI-like tool-calling path). What that backend actually needs is a - self-contained schema it can resolve without a document loader -- + self-contained schema it can resolve without a document loader - unlike a bare "$ref present" check, this asserts the referenced definition genuinely carries the two fields the rest of the pipeline (parse_ai_response, matching.py's resolve_*_ids) relies on. diff --git a/src/paperless_ai/tests/test_taxonomy.py b/src/paperless_ai/tests/test_taxonomy.py index 074c86075..470badc1c 100644 --- a/src/paperless_ai/tests/test_taxonomy.py +++ b/src/paperless_ai/tests/test_taxonomy.py @@ -24,13 +24,13 @@ class TestGetAssignedMetadata: GIVEN: - A document with no tags/type/correspondent/storage_path assigned WHEN: - - get_assigned_metadata() is called + - get_assigned_metadata() is called with no user (unrestricted) THEN: - All fields report as empty/None """ document = DocumentFactory.create() - result = get_assigned_metadata(document) + result = get_assigned_metadata(document, user=None) assert result == { "tags": [], @@ -44,7 +44,7 @@ class TestGetAssignedMetadata: GIVEN: - A document with tags, document_type, correspondent, and storage_path assigned WHEN: - - get_assigned_metadata() is called + - get_assigned_metadata() is called with no user (unrestricted) THEN: - All assigned fields are reported with their name values """ @@ -59,13 +59,78 @@ class TestGetAssignedMetadata: ) document.tags.add(tag) - result = get_assigned_metadata(document) + result = get_assigned_metadata(document, user=None) assert result["tags"] == ["Bloodwork"] assert result["document_type"] == "Lab Report" assert result["correspondent"] == "City Hospital" assert result["storage_path"] == "Medical" + def test_assigned_tag_invisible_to_user_is_omitted(self) -> None: + """ + GIVEN: + - A document with a tag owned by a different user + - A non-superuser requester with no visibility into that tag + WHEN: + - get_assigned_metadata() is called for the requester + THEN: + - The invisible tag's name is not surfaced - a document being + visible to a user does not imply every object assigned to it + is (per-object permissions can differ) + """ + tag_owner = UserFactory.create() + tag = TagFactory.create(name="Restricted", owner=tag_owner) + document = DocumentFactory.create() + document.tags.add(tag) + requester = UserFactory.create() + + result = get_assigned_metadata(document, user=requester) + + assert result["tags"] == [] + + def test_assigned_correspondent_invisible_to_user_is_omitted(self) -> None: + """ + GIVEN: + - A document whose correspondent is owned by a different user + - A non-superuser requester with no visibility into that + correspondent + WHEN: + - get_assigned_metadata() is called for the requester + THEN: + - The correspondent is reported as unset, not its actual name + """ + correspondent_owner = UserFactory.create() + correspondent = CorrespondentFactory.create( + name="Restricted Correspondent", + owner=correspondent_owner, + ) + document = DocumentFactory.create(correspondent=correspondent) + requester = UserFactory.create() + + result = get_assigned_metadata(document, user=requester) + + assert result["correspondent"] is None + + def test_assigned_metadata_visible_to_superuser(self) -> None: + """ + GIVEN: + - A document with a tag owned by a different user + - A superuser requester + WHEN: + - get_assigned_metadata() is called for the superuser + THEN: + - The tag's name is surfaced - superusers see everything + """ + tag_owner = UserFactory.create() + tag = TagFactory.create(name="Owned By Someone Else", owner=tag_owner) + document = DocumentFactory.create() + document.tags.add(tag) + superuser = UserFactory.create(is_superuser=True) + + result = get_assigned_metadata(document, user=superuser) + + assert result["tags"] == ["Owned By Someone Else"] + def make_node(document_id: int, score: float) -> SimpleNamespace: """A stand-in for NodeWithScore: only ``.metadata``/``.score`` are read.""" @@ -125,7 +190,7 @@ class TestBuildTaxonomyCandidates: THEN: - The candidate uses the current tag name, not the indexed name """ - # The node's own metadata name (if any) must never be trusted -- + # The node's own metadata name (if any) must never be trusted - # only the document_id is used to re-derive the current name. tag = TagFactory.create(name="Old Name") document = DocumentFactory.create()