From 84bd0d6fcdbf233d0291bcfbb8ef0c9e9d04eb8e Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Fri, 12 Jun 2026 19:46:49 -0700 Subject: [PATCH] Enhancement(beta): gate and assemble taxonomy hints for a document Co-Authored-By: Claude Opus 4.8 --- src/paperless_ai/taxonomy.py | 47 ++++++++++++- src/paperless_ai/tests/test_taxonomy.py | 94 +++++++++++++++++++++++++ 2 files changed, 138 insertions(+), 3 deletions(-) diff --git a/src/paperless_ai/taxonomy.py b/src/paperless_ai/taxonomy.py index dc320979c..6b9cb77d9 100644 --- a/src/paperless_ai/taxonomy.py +++ b/src/paperless_ai/taxonomy.py @@ -1,12 +1,16 @@ -import logging from typing import TYPE_CHECKING from typing import TypedDict +from django.contrib.auth.models import User + +from documents.models import Document +from documents.permissions import get_objects_for_user_owner_aware +from paperless.config import AIConfig +from paperless_ai.indexing import retrieve_similar_nodes + if TYPE_CHECKING: from llama_index.core.schema import NodeWithScore -logger = logging.getLogger("paperless_ai.taxonomy") - class TaxonomyHints(TypedDict): tags: list[str] @@ -87,3 +91,40 @@ def format_hints_for_prompt(hints: TaxonomyHints) -> str: return "" return "\n\n".join([*blocks, _HINT_INSTRUCTION]) + + +def get_taxonomy_hints_for_document( + document: Document, + user: User | None, +) -> TaxonomyHints | None: + """Build taxonomy hints from a document's RAG neighbours. + + Returns ``None`` when no embedding backend is configured (the gate) so the + caller's prompt and matching are identical to today. Otherwise returns a + ``TaxonomyHints`` -- possibly all-empty when no similar documents exist. + Applies the same owner-aware visible-document filter as + ``get_context_for_document``. + """ + if not AIConfig().llm_embedding_backend: + return None + + visible_documents = ( + get_objects_for_user_owner_aware( + user, + "view_document", + Document, + ) + if user + else None + ) + visible_document_ids = ( + list(visible_documents.values_list("pk", flat=True)) + if visible_documents is not None + else None + ) + + nodes = retrieve_similar_nodes( + document=document, + document_ids=visible_document_ids, + ) + return build_taxonomy_hints_from_nodes(nodes) diff --git a/src/paperless_ai/tests/test_taxonomy.py b/src/paperless_ai/tests/test_taxonomy.py index 0f4cbc45a..77e9da052 100644 --- a/src/paperless_ai/tests/test_taxonomy.py +++ b/src/paperless_ai/tests/test_taxonomy.py @@ -1,8 +1,12 @@ from types import SimpleNamespace +import pytest +import pytest_mock + from paperless_ai.taxonomy import TaxonomyHints from paperless_ai.taxonomy import build_taxonomy_hints_from_nodes from paperless_ai.taxonomy import format_hints_for_prompt +from paperless_ai.taxonomy import get_taxonomy_hints_for_document def make_node(**metadata: object) -> SimpleNamespace: @@ -127,3 +131,93 @@ class TestFormatHintsForPrompt: } result = format_hints_for_prompt(hints) assert result.count("Prefer existing names from these lists verbatim") == 1 + + +@pytest.mark.django_db +class TestGetTaxonomyHintsForDocument: + def test_returns_none_when_embedding_backend_off( + self, + mocker: pytest_mock.MockerFixture, + ) -> None: + mocker.patch( + "paperless_ai.taxonomy.AIConfig", + return_value=SimpleNamespace(llm_embedding_backend=None), + ) + retrieve = mocker.patch("paperless_ai.taxonomy.retrieve_similar_nodes") + + result = get_taxonomy_hints_for_document(SimpleNamespace(), user=None) + + assert result is None + retrieve.assert_not_called() + + def test_passes_owner_aware_ids_when_user_present( + self, + mocker: pytest_mock.MockerFixture, + ) -> None: + mocker.patch( + "paperless_ai.taxonomy.AIConfig", + return_value=SimpleNamespace(llm_embedding_backend="huggingface"), + ) + visible = mocker.MagicMock() + visible.values_list.return_value = [1, 2, 3] + mocker.patch( + "paperless_ai.taxonomy.get_objects_for_user_owner_aware", + return_value=visible, + ) + retrieve = mocker.patch( + "paperless_ai.taxonomy.retrieve_similar_nodes", + return_value=[], + ) + document = SimpleNamespace() + user = mocker.MagicMock() + + get_taxonomy_hints_for_document(document, user=user) + + retrieve.assert_called_once_with( + document=document, + document_ids=[1, 2, 3], + ) + + def test_returns_populated_hints_when_nodes_found( + self, + mocker: pytest_mock.MockerFixture, + ) -> None: + mocker.patch( + "paperless_ai.taxonomy.AIConfig", + return_value=SimpleNamespace(llm_embedding_backend="huggingface"), + ) + mocker.patch( + "paperless_ai.taxonomy.retrieve_similar_nodes", + return_value=[make_node(tags=["Taxes"], document_type="Invoice")], + ) + + result = get_taxonomy_hints_for_document(SimpleNamespace(), user=None) + + assert result == { + "tags": ["Taxes"], + "document_types": ["Invoice"], + "correspondents": [], + "storage_paths": [], + } + + def test_returns_empty_hints_not_none_when_no_nodes( + self, + mocker: pytest_mock.MockerFixture, + ) -> None: + mocker.patch( + "paperless_ai.taxonomy.AIConfig", + return_value=SimpleNamespace(llm_embedding_backend="huggingface"), + ) + mocker.patch( + "paperless_ai.taxonomy.retrieve_similar_nodes", + return_value=[], + ) + + result = get_taxonomy_hints_for_document(SimpleNamespace(), user=None) + + assert result == { + "tags": [], + "document_types": [], + "correspondents": [], + "storage_paths": [], + }