diff --git a/src/paperless_ai/ai_classifier.py b/src/paperless_ai/ai_classifier.py index 14735e419..c8c666316 100644 --- a/src/paperless_ai/ai_classifier.py +++ b/src/paperless_ai/ai_classifier.py @@ -5,13 +5,12 @@ from django.conf import settings from django.contrib.auth.models import User from documents.models import Document -from documents.permissions import get_objects_for_user_owner_aware +from documents.permissions import permitted_object_ids from paperless.config import AIConfig from paperless_ai.base_model import ClassificationSuggestions from paperless_ai.base_model import TaxonomyChoiceDict from paperless_ai.client import AIClient from paperless_ai.db import db_connection_released -from paperless_ai.indexing import _node_document_ids from paperless_ai.indexing import retrieve_similar_nodes from paperless_ai.indexing import truncate_content from paperless_ai.prompts.context import ClassificationPromptContext @@ -19,7 +18,9 @@ from paperless_ai.prompts.context import LocalizationPromptContext from paperless_ai.prompts.context import RagContextPromptContext from paperless_ai.prompts.render import render_prompt from paperless_ai.taxonomy import AssignedMetadata +from paperless_ai.taxonomy import SimilarDocument from paperless_ai.taxonomy import TaxonomyCandidates +from paperless_ai.taxonomy import _node_document_weights from paperless_ai.taxonomy import build_taxonomy_candidates from paperless_ai.taxonomy import empty_taxonomy_candidates from paperless_ai.taxonomy import format_taxonomy_for_prompt @@ -39,6 +40,35 @@ logger = logging.getLogger("paperless_ai.rag_classifier") TAXONOMY_CANDIDATE_TOP_K = 15 +def _fulltext_similar_documents( + document: Document, + user: User | None, + top_k: int, +) -> list[SimilarDocument]: + """Rank-based fallback when no embedding backend is configured. Uses + Tantivy's "More Like This" (term-overlap similarity) instead of vector + similarity - cruder, but far better than no candidates at all. + more_like_this_ids returns only a ranked ID list, no scores, so weight is + synthesized from rank (descending from top_k) rather than claiming a + similarity magnitude that doesn't exist. A superuser is normalized to + ``None`` before calling, since the backend's permission filter has no + superuser short-circuit of its own. + """ + from documents.search import get_backend + + search_user = None if user is not None and user.is_superuser else user + backend = get_backend() + similar_ids = backend.more_like_this_ids( + document.pk, + user=search_user, + limit=top_k, + ) + return [ + SimilarDocument(document_id=doc_id, weight=float(top_k - rank)) + for rank, doc_id in enumerate(similar_ids) + ] + + def get_language_name(language_code: str) -> str: normalized_language_code = language_code.lower() for code, name in settings.LANGUAGES: @@ -131,44 +161,52 @@ def get_taxonomy_context( user: User | None = None, max_docs: int = 5, ) -> tuple[TaxonomyCandidates, AssignedMetadata, str]: - """One retrieval feeds both taxonomy candidates and RAG text context. - On any retrieval failure, degrades to empty candidates/context rather than - propagating the exception - a vector-store outage should not block - classification, only its RAG-assisted enrichment. + """One retrieval feeds both taxonomy candidates and RAG text context. Uses + vector similarity when an embedding backend is configured, otherwise + falls back to Tantivy full-text "More Like This" similarity - see + _fulltext_similar_documents. On any retrieval failure, degrades to empty + candidates/context rather than propagating the exception - neither a + vector-store outage nor a search-index issue should block classification, + only its context-assisted enrichment. """ assigned = get_assigned_metadata(document, user) + ai_config = AIConfig() try: - # None means "no restriction" to retrieve_similar_nodes. A superuser - # (like no user at all) can see every document, so skip materializing - # every visible pk into a Python list and passing it through as an IN - # filter: for a large library that is a wasted quadratic scan in the - # vector store at best, and past ~32,763 documents a hard - # sqlite3.OperationalError (SQLite's bound-parameter limit) at worst. - # get_objects_for_user_owner_aware() would return every Document for a - # superuser anyway (guardian's own with_superuser shortcut), so this - # changes nothing about which documents are considered -- only how we - # get there. - visible_document_ids = ( - None - if user is None or user.is_superuser - else list( - get_objects_for_user_owner_aware( - user, - "view_document", - Document, - ).values_list("pk", flat=True), + if ai_config.llm_embedding_backend: + # None means "no restriction" to retrieve_similar_nodes. A superuser + # (like no user at all) can see every document, so skip materializing + # every visible pk into a Python list and passing it through as an IN + # filter: for a large library that is a wasted quadratic scan in the + # vector store at best, and past ~32,763 documents a hard + # sqlite3.OperationalError (SQLite's bound-parameter limit) at worst. + # permitted_object_ids() has its own superuser shortcut that would + # return every Document's id anyway, so this changes nothing about + # which documents are considered -- only how we get there. + visible_document_ids = ( + None + if user is None or user.is_superuser + else list(permitted_object_ids(user, Document, "view_document")) + ) + nodes = retrieve_similar_nodes( + document, + top_k=TAXONOMY_CANDIDATE_TOP_K, + document_ids=visible_document_ids, + ) + similar_documents = _node_document_weights(nodes) + else: + # See _fulltext_similar_documents: it applies its own permission + # filter via `user`, so no visible-document-id list is needed here. + similar_documents = _fulltext_similar_documents( + document, + user, + top_k=TAXONOMY_CANDIDATE_TOP_K, ) - ) - nodes = retrieve_similar_nodes( - document, - top_k=TAXONOMY_CANDIDATE_TOP_K, - document_ids=visible_document_ids, - ) - candidates = build_taxonomy_candidates(nodes, user) + candidates = build_taxonomy_candidates(similar_documents, user) + similar_doc_ids = [s["document_id"] for s in similar_documents] similar_docs = list( - Document.objects.filter(pk__in=_node_document_ids(nodes))[:max_docs], + Document.objects.filter(pk__in=similar_doc_ids)[:max_docs], ) context_blocks = [] for similar in similar_docs: @@ -177,8 +215,8 @@ def get_taxonomy_context( context_blocks.append(f"TITLE: {title}\n{text}") except Exception: logger.exception( - "Failed to retrieve RAG neighbours for document %s; continuing " - "without taxonomy candidates or similar-document context.", + "Failed to retrieve similar-document context for document %s; " + "continuing without taxonomy candidates or similar-document context.", document.pk, ) return empty_taxonomy_candidates(), assigned, "" @@ -261,23 +299,14 @@ def get_ai_document_classification( ) -> ClassificationSuggestions: ai_config = AIConfig() - if ai_config.llm_embedding_backend: - candidates, assigned, context = get_taxonomy_context(document, user) - prompt = build_prompt_with_rag( - document, - ai_config, - candidates=candidates, - assigned=assigned, - context=context, - ) - else: - candidates = empty_taxonomy_candidates() - prompt = build_prompt_without_rag( - document, - ai_config, - candidates=candidates, - assigned=get_assigned_metadata(document, user), - ) + candidates, assigned, context = get_taxonomy_context(document, user) + prompt = build_prompt_with_rag( + document, + ai_config, + candidates=candidates, + assigned=assigned, + context=context, + ) client = AIClient() # Hand the pooled DB connection back while the (slow) LLM query runs so it diff --git a/src/paperless_ai/taxonomy.py b/src/paperless_ai/taxonomy.py index 363242c35..dedf0fdba 100644 --- a/src/paperless_ai/taxonomy.py +++ b/src/paperless_ai/taxonomy.py @@ -33,6 +33,11 @@ class TaxonomyCandidate(TypedDict): weight: float +class SimilarDocument(TypedDict): + document_id: int + weight: float + + class TaxonomyCandidates(TypedDict): tags: list[TaxonomyCandidate] document_types: list[TaxonomyCandidate] @@ -105,10 +110,10 @@ def get_assigned_metadata(document: Document, user: User | None) -> AssignedMeta ) -def _node_document_weights(nodes: list["NodeWithScore"]) -> dict[int, float]: - """document_id -> that node's similarity score, summed if a document_id - appears more than once across the retrieved nodes (e.g. multiple chunks - of the same source document).""" +def _node_document_weights(nodes: list["NodeWithScore"]) -> list[SimilarDocument]: + """Sum each node's similarity score into its document_id (a document can + appear via multiple chunks/nodes) and return one SimilarDocument per + distinct document_id.""" weights: dict[int, float] = defaultdict(float) for node in nodes: document_id = node.metadata.get("document_id") @@ -121,7 +126,10 @@ def _node_document_weights(nodes: list["NodeWithScore"]) -> dict[int, float]: weights[int(document_id)] += float(node.score or 0.0) except (TypeError, ValueError): # pragma: no cover continue - return weights + return [ + SimilarDocument(document_id=document_id, weight=weight) + for document_id, weight in weights.items() + ] def _visible_ranked_candidates( @@ -157,21 +165,26 @@ def _visible_ranked_candidates( def build_taxonomy_candidates( - nodes: list["NodeWithScore"], + similar_documents: list[SimilarDocument], user: User | None, ) -> TaxonomyCandidates: - """Resolve each neighbour node's document_id to a live Document, read its - *current* tags/type/correspondent/storage_path via the ORM (never the - possibly-stale names cached in vector-index node metadata), weight each - distinct taxonomy object by aggregate neighbour similarity, permission-filter + """Resolve each similar document's id to a live Document, read its + *current* tags/type/correspondent/storage_path via the ORM (never any + possibly-stale names an adapter's source might have cached), weight each + distinct taxonomy object by aggregate similarity weight, permission-filter against what ``user`` can see, and return each category ranked by weight - and capped. + and capped. ``similar_documents`` may come from either the vector-RAG + adapter or the full-text fallback adapter - both produce this same shape. """ - - document_weights = _node_document_weights(nodes) - if not document_weights: + if not similar_documents: return empty_taxonomy_candidates() + # Both adapters guarantee at most one SimilarDocument per document_id, so + # this never silently drops a duplicate's weight. + document_weights: dict[int, float] = { + s["document_id"]: s["weight"] for s in similar_documents + } + # Only .tags.all() needs prefetching (a reverse M2M, one extra query for # the whole batch). document_type/correspondent/storage_path are read # below via their *_id columns (neighbour.document_type_id, etc.), which diff --git a/src/paperless_ai/tests/test_ai_classifier.py b/src/paperless_ai/tests/test_ai_classifier.py index 6159f78b9..289beb8c5 100644 --- a/src/paperless_ai/tests/test_ai_classifier.py +++ b/src/paperless_ai/tests/test_ai_classifier.py @@ -1,3 +1,4 @@ +from collections.abc import Generator from types import SimpleNamespace from unittest.mock import MagicMock from unittest.mock import patch @@ -7,10 +8,13 @@ import pytest_mock from django.test import override_settings from documents.models import Document +from documents.search import TantivyBackend 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 TAXONOMY_CANDIDATE_TOP_K +from paperless_ai.ai_classifier import _fulltext_similar_documents 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 @@ -20,6 +24,7 @@ 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 SimilarDocument from paperless_ai.taxonomy import TaxonomyCandidate from paperless_ai.taxonomy import TaxonomyCandidates from paperless_ai.taxonomy import empty_taxonomy_candidates @@ -204,12 +209,10 @@ def test_use_rag_if_configured( @pytest.mark.django_db @patch("paperless_ai.client.AIClient.run_llm_query") -@patch("paperless_ai.ai_classifier.build_prompt_without_rag") -@patch("paperless_ai.ai_classifier.AIConfig") +@patch("paperless_ai.ai_classifier.build_prompt_with_rag") @override_settings(LLM_BACKEND="ollama", LLM_MODEL="some_model") -def test_use_without_rag_if_not_configured( - mock_ai_config, - mock_build_prompt_without_rag, +def test_use_rag_prompt_even_without_embedding_backend( + mock_build_prompt_with_rag, mock_run_llm_query, mock_document, ): @@ -219,13 +222,13 @@ def test_use_without_rag_if_not_configured( WHEN: - get_ai_document_classification() is called THEN: - - The non-RAG prompt builder is used + - The RAG-context prompt builder is still used (fed by the full-text + fallback's context/candidates instead of the vector store's) """ - mock_ai_config.return_value.llm_embedding_backend = None - mock_build_prompt_without_rag.return_value = "Prompt without RAG" + mock_build_prompt_with_rag.return_value = "Prompt with RAG" mock_run_llm_query.return_value = NESTED_SUGGESTIONS get_ai_document_classification(mock_document) - mock_build_prompt_without_rag.assert_called_once() + mock_build_prompt_with_rag.assert_called_once() @pytest.mark.django_db @@ -303,6 +306,7 @@ def test_build_localization_prompt_preserves_unicode_characters(): @pytest.mark.django_db +@override_settings(LLM_EMBEDDING_BACKEND="huggingface") def test_get_taxonomy_context_assembles_rag_text_and_candidates(): """ GIVEN: @@ -344,6 +348,7 @@ def test_get_taxonomy_context_assembles_rag_text_and_candidates(): @pytest.mark.django_db +@override_settings(LLM_EMBEDDING_BACKEND="huggingface") def test_get_taxonomy_context_no_similar_docs(): """ GIVEN: @@ -367,6 +372,67 @@ def test_get_taxonomy_context_no_similar_docs(): } +@pytest.mark.django_db +def test_get_taxonomy_context_uses_fulltext_fallback_when_no_embedding_backend( + mocker: pytest_mock.MockerFixture, +) -> None: + """ + GIVEN: + - No LLM embedding backend is configured (the default test settings) + WHEN: + - get_taxonomy_context() is called + THEN: + - _fulltext_similar_documents() is called with the document, the user + and TAXONOMY_CANDIDATE_TOP_K + - retrieve_similar_nodes() (the vector path) is never called + """ + document = DocumentFactory.create(content="Some content") + mock_fulltext = mocker.patch( + "paperless_ai.ai_classifier._fulltext_similar_documents", + return_value=[], + ) + mock_retrieve = mocker.patch("paperless_ai.ai_classifier.retrieve_similar_nodes") + + get_taxonomy_context(document, user=None) + + mock_fulltext.assert_called_once_with( + document, + None, + top_k=TAXONOMY_CANDIDATE_TOP_K, + ) + mock_retrieve.assert_not_called() + + +@pytest.mark.django_db +@override_settings(LLM_EMBEDDING_BACKEND="huggingface") +def test_get_taxonomy_context_uses_vector_path_when_embedding_backend_configured( + mocker: pytest_mock.MockerFixture, +) -> None: + """ + GIVEN: + - An LLM embedding backend is configured + WHEN: + - get_taxonomy_context() is called + THEN: + - retrieve_similar_nodes() (the vector path) is called + - _fulltext_similar_documents() (the no-embedding-backend fallback) + is never called + """ + document = DocumentFactory.create(content="Some content") + mock_retrieve = mocker.patch( + "paperless_ai.ai_classifier.retrieve_similar_nodes", + return_value=[], + ) + mock_fulltext = mocker.patch( + "paperless_ai.ai_classifier._fulltext_similar_documents", + ) + + get_taxonomy_context(document, user=None) + + mock_retrieve.assert_called_once() + mock_fulltext.assert_not_called() + + class TestGetTaxonomyContextVisibility: """get_taxonomy_context must not materialize every visible document id for a user who can already see the whole library: a superuser (like no @@ -379,6 +445,7 @@ class TestGetTaxonomyContextVisibility: """ @pytest.mark.django_db + @override_settings(LLM_EMBEDDING_BACKEND="huggingface") def test_skips_permission_lookup_for_superuser( self, mocker: pytest_mock.MockerFixture, @@ -397,17 +464,18 @@ class TestGetTaxonomyContextVisibility: "paperless_ai.ai_classifier.retrieve_similar_nodes", return_value=[], ) - mock_get_objects = mocker.patch( - "paperless_ai.ai_classifier.get_objects_for_user_owner_aware", + mock_permitted = mocker.patch( + "paperless_ai.ai_classifier.permitted_object_ids", ) user = UserFactory.create(is_superuser=True) get_taxonomy_context(document, user) - mock_get_objects.assert_not_called() + mock_permitted.assert_not_called() assert mock_retrieve.call_args.kwargs["document_ids"] is None @pytest.mark.django_db + @override_settings(LLM_EMBEDDING_BACKEND="huggingface") def test_skips_permission_lookup_when_no_user( self, mocker: pytest_mock.MockerFixture, @@ -426,16 +494,17 @@ class TestGetTaxonomyContextVisibility: "paperless_ai.ai_classifier.retrieve_similar_nodes", return_value=[], ) - mock_get_objects = mocker.patch( - "paperless_ai.ai_classifier.get_objects_for_user_owner_aware", + mock_permitted = mocker.patch( + "paperless_ai.ai_classifier.permitted_object_ids", ) get_taxonomy_context(document, None) - mock_get_objects.assert_not_called() + mock_permitted.assert_not_called() assert mock_retrieve.call_args.kwargs["document_ids"] is None @pytest.mark.django_db + @override_settings(LLM_EMBEDDING_BACKEND="huggingface") def test_restricts_to_visible_documents_for_non_superuser( self, mocker: pytest_mock.MockerFixture, @@ -446,7 +515,7 @@ class TestGetTaxonomyContextVisibility: WHEN: - get_taxonomy_context() is called THEN: - - The user's visible document ids are looked up and passed to + - The user's permitted document ids are looked up and passed to retrieve_similar_nodes() as a restriction """ document = DocumentFactory.create(content="Some content") @@ -454,21 +523,185 @@ class TestGetTaxonomyContextVisibility: "paperless_ai.ai_classifier.retrieve_similar_nodes", return_value=[], ) - mock_queryset = mocker.MagicMock() - mock_queryset.values_list.return_value = [1, 2, 3] - mock_get_objects = mocker.patch( - "paperless_ai.ai_classifier.get_objects_for_user_owner_aware", - return_value=mock_queryset, + mock_permitted = mocker.patch( + "paperless_ai.ai_classifier.permitted_object_ids", + return_value=[1, 2, 3], ) user = UserFactory.create(is_superuser=False) get_taxonomy_context(document, user) - mock_get_objects.assert_called_once_with(user, "view_document", Document) + mock_permitted.assert_called_once_with(user, Document, "view_document") assert mock_retrieve.call_args.kwargs["document_ids"] == [1, 2, 3] @pytest.mark.django_db +class TestFulltextSimilarDocuments: + """_fulltext_similar_documents is the no-embedding-backend fallback: it + asks the Tantivy full-text index for "More Like This" neighbours instead + of the vector store, and synthesizes a rank-based weight since Tantivy's + more_like_this_ids returns only an ordered id list, no scores. + """ + + @pytest.fixture + def fulltext_backend( + self, + mocker: pytest_mock.MockerFixture, + ) -> Generator[TantivyBackend, None, None]: + """An in-memory Tantivy backend, wired up as the module-level + singleton _fulltext_similar_documents resolves via get_backend().""" + backend = TantivyBackend(path=None) + backend.open() + mocker.patch("documents.search.get_backend", return_value=backend) + try: + yield backend + finally: + backend.close() + + def test_ranks_by_rank_based_weight_descending( + self, + fulltext_backend: TantivyBackend, + ) -> None: + """ + GIVEN: + - A source document and two similar documents indexed in Tantivy + WHEN: + - _fulltext_similar_documents() is called + THEN: + - Each result's weight reflects its rank (first result weighted + higher than the second), not a raw similarity score + """ + source = DocumentFactory.create(content="quarterly financial report details") + first = DocumentFactory.create(content="quarterly financial report details") + second = DocumentFactory.create(content="financial report") + for doc in (source, first, second): + fulltext_backend.add_or_update(doc) + + result = _fulltext_similar_documents(source, user=None, top_k=5) + + assert len(result) == 2 + weight_by_id = {s["document_id"]: s["weight"] for s in result} + assert weight_by_id[first.pk] > weight_by_id[second.pk] + + def test_excludes_source_document( + self, + fulltext_backend: TantivyBackend, + ) -> None: + """ + GIVEN: + - A source document indexed in Tantivy with no other documents + WHEN: + - _fulltext_similar_documents() is called + THEN: + - An empty list is returned - the source document is never its + own similar document + """ + source = DocumentFactory.create(content="unique unrelated content") + fulltext_backend.add_or_update(source) + + result = _fulltext_similar_documents(source, user=None, top_k=5) + + assert result == [] + + def test_empty_index_returns_empty_list( + self, + fulltext_backend: TantivyBackend, + ) -> None: + """ + GIVEN: + - A document that has never been indexed (fresh/empty Tantivy index) + WHEN: + - _fulltext_similar_documents() is called + THEN: + - An empty list is returned rather than raising + """ + source = DocumentFactory.create(content="never indexed") + + result = _fulltext_similar_documents(source, user=None, top_k=5) + + assert result == [] + + def test_respects_top_k_limit( + self, + fulltext_backend: TantivyBackend, + ) -> None: + """ + GIVEN: + - A source document and four similar documents indexed + WHEN: + - _fulltext_similar_documents() is called with top_k=2 + THEN: + - At most 2 results are returned + """ + source = DocumentFactory.create(content="shared overlapping keyword text") + for _ in range(4): + fulltext_backend.add_or_update( + DocumentFactory.create(content="shared overlapping keyword text"), + ) + + result = _fulltext_similar_documents(source, user=None, top_k=2) + + assert len(result) <= 2 + + def test_result_shape_is_similar_document( + self, + fulltext_backend: TantivyBackend, + ) -> None: + """ + GIVEN: + - A source document and one similar document indexed + WHEN: + - _fulltext_similar_documents() is called + THEN: + - Each result is a SimilarDocument (document_id + weight only) + """ + source = DocumentFactory.create(content="shared content phrase") + other = DocumentFactory.create(content="shared content phrase") + fulltext_backend.add_or_update(source) + fulltext_backend.add_or_update(other) + + result = _fulltext_similar_documents(source, user=None, top_k=5) + + # rank 0 (the only/best result) with top_k=5 -> weight = top_k - rank = 5.0, + # per the "first result gets top_k, the last gets 1" formula. + assert result == [SimilarDocument(document_id=other.pk, weight=5.0)] + + def test_superuser_sees_other_users_documents( + self, + fulltext_backend: TantivyBackend, + ) -> None: + """ + GIVEN: + - A source document owned by one user and a similar document + owned by a different user, with no sharing between them + WHEN: + - _fulltext_similar_documents() is called with a superuser + THEN: + - The other user's document is still returned as a similar + document - a superuser must not be narrowed by the backend's + owner-based permission filter + """ + owner = UserFactory.create() + other_owner = UserFactory.create() + superuser = UserFactory.create(is_superuser=True) + source = DocumentFactory.create( + content="shared content phrase", + owner=owner, + ) + other = DocumentFactory.create( + content="shared content phrase", + owner=other_owner, + ) + fulltext_backend.add_or_update(source) + fulltext_backend.add_or_update(other) + + result = _fulltext_similar_documents(source, user=superuser, top_k=5) + + assert [s["document_id"] for s in result] == [other.pk] + + +@pytest.mark.django_db +@override_settings(LLM_EMBEDDING_BACKEND="huggingface") @patch("paperless_ai.ai_classifier.retrieve_similar_nodes") def test_get_taxonomy_context_retrieval_failure_degrades_to_no_hints(mock_retrieve): """ @@ -495,6 +728,7 @@ def test_get_taxonomy_context_retrieval_failure_degrades_to_no_hints(mock_retrie @pytest.mark.django_db +@override_settings(LLM_EMBEDDING_BACKEND="huggingface") @patch("paperless_ai.ai_classifier.build_taxonomy_candidates") @patch("paperless_ai.ai_classifier.retrieve_similar_nodes") def test_get_taxonomy_context_candidate_building_failure_degrades_to_no_hints( diff --git a/src/paperless_ai/tests/test_taxonomy.py b/src/paperless_ai/tests/test_taxonomy.py index 26bcab7bc..088cdf8ed 100644 --- a/src/paperless_ai/tests/test_taxonomy.py +++ b/src/paperless_ai/tests/test_taxonomy.py @@ -1,5 +1,4 @@ import json -from types import SimpleNamespace import pytest import pytest_mock @@ -11,6 +10,7 @@ from documents.tests.factories import StoragePathFactory from documents.tests.factories import TagFactory from documents.tests.factories import UserFactory from paperless_ai.taxonomy import AssignedMetadata +from paperless_ai.taxonomy import SimilarDocument from paperless_ai.taxonomy import TaxonomyCandidates from paperless_ai.taxonomy import build_taxonomy_candidates from paperless_ai.taxonomy import format_taxonomy_for_prompt @@ -132,9 +132,8 @@ class TestGetAssignedMetadata: 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.""" - return SimpleNamespace(metadata={"document_id": str(document_id)}, score=score) +def make_similar(document_id: int, weight: float) -> SimilarDocument: + return SimilarDocument(document_id=document_id, weight=weight) @pytest.mark.django_db @@ -170,9 +169,9 @@ class TestBuildTaxonomyCandidates: doc_a.tags.add(tag) doc_b = DocumentFactory.create() doc_b.tags.add(tag) - nodes = [make_node(doc_a.pk, 0.9), make_node(doc_b.pk, 0.4)] + similar_documents = [make_similar(doc_a.pk, 0.9), make_similar(doc_b.pk, 0.4)] - result = build_taxonomy_candidates(nodes, user=None) + result = build_taxonomy_candidates(similar_documents, user=None) assert len(result["tags"]) == 1 assert result["tags"][0]["id"] == tag.pk @@ -197,9 +196,9 @@ class TestBuildTaxonomyCandidates: document.tags.add(tag) tag.name = "New Name" tag.save() - nodes = [make_node(document.pk, 0.5)] + similar_documents = [make_similar(document.pk, 0.5)] - result = build_taxonomy_candidates(nodes, user=None) + result = build_taxonomy_candidates(similar_documents, user=None) assert result["tags"][0]["name"] == "New Name" @@ -219,9 +218,9 @@ class TestBuildTaxonomyCandidates: document = DocumentFactory.create() document.tags.add(tag) tag.delete() - nodes = [make_node(document.pk, 0.5)] + similar_documents = [make_similar(document.pk, 0.5)] - result = build_taxonomy_candidates(nodes, user=None) + result = build_taxonomy_candidates(similar_documents, user=None) assert result["tags"] == [] @@ -240,9 +239,12 @@ class TestBuildTaxonomyCandidates: strong_doc.tags.add(strong_tag) weak_doc = DocumentFactory.create() weak_doc.tags.add(weak_tag) - nodes = [make_node(strong_doc.pk, 0.9), make_node(weak_doc.pk, 0.1)] + similar_documents = [ + make_similar(strong_doc.pk, 0.9), + make_similar(weak_doc.pk, 0.1), + ] - result = build_taxonomy_candidates(nodes, user=None) + result = build_taxonomy_candidates(similar_documents, user=None) assert [c["name"] for c in result["tags"]] == ["Strong", "Weak"] @@ -258,9 +260,9 @@ class TestBuildTaxonomyCandidates: document = DocumentFactory.create() for i in range(15): document.tags.add(TagFactory.create(name=f"Tag{i}")) - nodes = [make_node(document.pk, 0.5)] + similar_documents = [make_similar(document.pk, 0.5)] - result = build_taxonomy_candidates(nodes, user=None) + result = build_taxonomy_candidates(similar_documents, user=None) assert len(result["tags"]) == 10 @@ -274,12 +276,12 @@ class TestBuildTaxonomyCandidates: - Only 5 correspondents are returned """ correspondents = CorrespondentFactory.create_batch(7) - nodes = [ - make_node(DocumentFactory.create(correspondent=c).pk, 0.5) + similar_documents = [ + make_similar(DocumentFactory.create(correspondent=c).pk, 0.5) for c in correspondents ] - result = build_taxonomy_candidates(nodes, user=None) + result = build_taxonomy_candidates(similar_documents, user=None) assert len(result["correspondents"]) == 5 @@ -294,9 +296,9 @@ class TestBuildTaxonomyCandidates: """ document_type = DocumentTypeFactory.create(name="Invoice") document = DocumentFactory.create(document_type=document_type) - nodes = [make_node(document.pk, 0.5)] + similar_documents = [make_similar(document.pk, 0.5)] - result = build_taxonomy_candidates(nodes, user=None) + result = build_taxonomy_candidates(similar_documents, user=None) assert len(result["document_types"]) == 1 assert result["document_types"][0]["id"] == document_type.pk @@ -312,12 +314,12 @@ class TestBuildTaxonomyCandidates: - Only 5 document_types are returned """ document_types = DocumentTypeFactory.create_batch(7) - nodes = [ - make_node(DocumentFactory.create(document_type=dt).pk, 0.5) + similar_documents = [ + make_similar(DocumentFactory.create(document_type=dt).pk, 0.5) for dt in document_types ] - result = build_taxonomy_candidates(nodes, user=None) + result = build_taxonomy_candidates(similar_documents, user=None) assert len(result["document_types"]) == 5 @@ -332,9 +334,9 @@ class TestBuildTaxonomyCandidates: """ storage_path = StoragePathFactory.create(name="Invoices") document = DocumentFactory.create(storage_path=storage_path) - nodes = [make_node(document.pk, 0.5)] + similar_documents = [make_similar(document.pk, 0.5)] - result = build_taxonomy_candidates(nodes, user=None) + result = build_taxonomy_candidates(similar_documents, user=None) assert len(result["storage_paths"]) == 1 assert result["storage_paths"][0]["id"] == storage_path.pk @@ -350,12 +352,12 @@ class TestBuildTaxonomyCandidates: - Only 5 storage_paths are returned """ storage_paths = StoragePathFactory.create_batch(7) - nodes = [ - make_node(DocumentFactory.create(storage_path=sp).pk, 0.5) + similar_documents = [ + make_similar(DocumentFactory.create(storage_path=sp).pk, 0.5) for sp in storage_paths ] - result = build_taxonomy_candidates(nodes, user=None) + result = build_taxonomy_candidates(similar_documents, user=None) assert len(result["storage_paths"]) == 5 @@ -375,14 +377,14 @@ class TestBuildTaxonomyCandidates: tag = TagFactory.create(name="Restricted") document = DocumentFactory.create() document.tags.add(tag) - nodes = [make_node(document.pk, 0.5)] + similar_documents = [make_similar(document.pk, 0.5)] user = UserFactory.create() mocker.patch( "documents.permissions.permitted_object_ids", return_value=[], # user cannot see this tag ) - result = build_taxonomy_candidates(nodes, user=user) + result = build_taxonomy_candidates(similar_documents, user=user) assert result["tags"] == [] @@ -412,10 +414,10 @@ class TestBuildTaxonomyCandidates: tag.save() document = DocumentFactory.create() document.tags.add(tag) - nodes = [make_node(document.pk, 0.5)] + similar_documents = [make_similar(document.pk, 0.5)] spy = mocker.patch("documents.permissions.permitted_object_ids") - result = build_taxonomy_candidates(nodes, user=None) + result = build_taxonomy_candidates(similar_documents, user=None) assert result["tags"][0]["name"] == "Owned" spy.assert_not_called()