from types import SimpleNamespace from unittest.mock import MagicMock from unittest.mock import patch import pytest import pytest_mock from django.test import override_settings from documents.models import Document 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 def mock_document(): doc = MagicMock(spec=Document) doc.title = "Test Title" doc.filename = "test_file.pdf" doc.created = "2023-01-01" doc.added = "2023-01-02" doc.modified = "2023-01-03" tag1 = MagicMock() tag1.name = "Tag1" tag2 = MagicMock() tag2.name = "Tag2" doc.tags.all = MagicMock(return_value=[tag1, tag2]) doc.document_type = MagicMock() doc.document_type.name = "Invoice" doc.correspondent = MagicMock() doc.correspondent.name = "Test Correspondent" doc.storage_path = None # get_assigned_metadata reads this directly doc.archive_serial_number = "12345" doc.content = "This is the document content." cf1 = MagicMock(__str__=lambda x: "Value1") cf1.field = MagicMock() cf1.field.name = "Field1" cf1.value = "Value1" cf2 = MagicMock(__str__=lambda x: "Value2") cf2.field = MagicMock() cf2.field.name = "Field2" cf2.value = "Value2" doc.custom_fields.all = MagicMock(return_value=[cf1, cf2]) return doc NESTED_SUGGESTIONS = { "title": "Test Title", "tags": {"existing_ids": [], "new_names": ["test", "document"]}, "correspondents": {"existing_ids": [], "new_names": ["John Doe"]}, "document_types": {"existing_ids": [], "new_names": ["report"]}, "storage_paths": {"existing_ids": [], "new_names": ["Reports"]}, "dates": ["2023-01-01"], } @pytest.mark.django_db @patch("paperless_ai.client.AIClient.run_llm_query") @override_settings(LLM_BACKEND="ollama", LLM_MODEL="some_model") def test_get_ai_document_classification_success(mock_run_llm_query, mock_document): """ GIVEN: - An LLM backend configured without RAG - A classification call followed by a localization call WHEN: - get_ai_document_classification() is called with an output_language THEN: - The localized title/new_names are used - Correspondents are never localized, so the original suggestion survives - Dates are never localized - The classification prompt has no taxonomy title instruction and the localization prompt asks to rewrite only new_names/title """ mock_run_llm_query.side_effect = [ NESTED_SUGGESTIONS, { "title": "Testtitel", "tags": {"existing_ids": [], "new_names": ["Test", "Document"]}, "correspondents": {"existing_ids": [], "new_names": ["Jane Doe"]}, "document_types": {"existing_ids": [], "new_names": ["Bericht"]}, "storage_paths": {"existing_ids": [], "new_names": ["Berichte"]}, "dates": ["2024-01-01"], }, ] result = get_ai_document_classification(mock_document, output_language="de-de") assert result["title"] == "Testtitel" assert result["tags"]["new_names"] == ["Test", "Document"] # Correspondents are never localized - the merge step doesn't touch them, # so the original (English) suggestion survives, same as before this change. assert result["correspondents"]["new_names"] == ["John Doe"] assert result["document_types"]["new_names"] == ["Bericht"] assert result["storage_paths"]["new_names"] == ["Berichte"] assert result["dates"] == ["2023-01-01"] classification_prompt = mock_run_llm_query.call_args_list[0].args[0] localization_prompt = mock_run_llm_query.call_args_list[1].args[0] assert "Write suggested titles" not in classification_prompt assert "Rewrite only the" in localization_prompt assert "Do not translate correspondents or dates" in localization_prompt @pytest.mark.django_db @patch("paperless_ai.client.AIClient.run_llm_query") @override_settings(LLM_BACKEND="ollama", LLM_MODEL="some_model") def test_get_ai_document_classification_keeps_originals_when_localization_empty( mock_run_llm_query, mock_document, ): """ GIVEN: - A localization response whose fields are all empty WHEN: - get_ai_document_classification() is called with an output_language THEN: - The original (pre-localization) suggestions are kept for every field """ mock_run_llm_query.side_effect = [ NESTED_SUGGESTIONS, { "title": "", "tags": {"existing_ids": [], "new_names": []}, "correspondents": {"existing_ids": [], "new_names": []}, "document_types": {"existing_ids": [], "new_names": []}, "storage_paths": {"existing_ids": [], "new_names": []}, "dates": [], }, ] result = get_ai_document_classification(mock_document, output_language="de-de") assert result["title"] == "Test Title" assert result["tags"]["new_names"] == ["test", "document"] assert result["correspondents"]["new_names"] == ["John Doe"] assert result["document_types"]["new_names"] == ["report"] assert result["storage_paths"]["new_names"] == ["Reports"] assert result["dates"] == ["2023-01-01"] @pytest.mark.django_db @patch("paperless_ai.client.AIClient.run_llm_query") def test_get_ai_document_classification_failure(mock_run_llm_query, mock_document): """ GIVEN: - The LLM client raises an exception WHEN: - get_ai_document_classification() is called THEN: - The exception propagates rather than being swallowed """ mock_run_llm_query.side_effect = Exception("LLM query failed") with pytest.raises(Exception): get_ai_document_classification(mock_document) @pytest.mark.django_db @patch("paperless_ai.client.AIClient.run_llm_query") @patch("paperless_ai.ai_classifier.build_prompt_with_rag") @patch("paperless_ai.ai_classifier.retrieve_similar_nodes") @override_settings( LLM_EMBEDDING_BACKEND="huggingface", LLM_EMBEDDING_MODEL="some_model", LLM_BACKEND="ollama", LLM_MODEL="some_model", ) def test_use_rag_if_configured( mock_retrieve, mock_build_prompt_with_rag, mock_run_llm_query, mock_document, ): """ GIVEN: - An LLM embedding backend is configured WHEN: - get_ai_document_classification() is called THEN: - The RAG-augmented prompt builder is used """ mock_retrieve.return_value = [] 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_with_rag.assert_called_once() @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") @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, mock_run_llm_query, mock_document, ): """ GIVEN: - No LLM embedding backend is configured WHEN: - get_ai_document_classification() is called THEN: - The non-RAG prompt builder is used """ mock_ai_config.return_value.llm_embedding_backend = None mock_build_prompt_without_rag.return_value = "Prompt without RAG" mock_run_llm_query.return_value = NESTED_SUGGESTIONS get_ai_document_classification(mock_document) mock_build_prompt_without_rag.assert_called_once() @pytest.mark.django_db @override_settings( LLM_EMBEDDING_BACKEND="huggingface", LLM_BACKEND="ollama", LLM_MODEL="some_model", ) def test_prompt_with_without_rag(mock_document): """ GIVEN: - A document and an AIConfig WHEN: - build_prompt_without_rag(), build_prompt_with_rag(), and build_localization_prompt() are called THEN: - build_prompt_without_rag() has no similar-documents section - build_prompt_with_rag() includes the similar-documents context - build_localization_prompt() asks to rewrite only new_names/title and not to translate correspondents or dates """ config = AIConfig() prompt = build_prompt_without_rag(mock_document, config) assert "Additional context from similar documents" not in prompt assert "for generated" not in prompt prompt = build_prompt_with_rag( mock_document, config, context="Context from similar documents", ) assert "Additional context from similar documents" in prompt assert "Context from similar documents" in prompt prompt = build_localization_prompt(NESTED_SUGGESTIONS, output_language="de-de") assert "Rewrite only the" in prompt assert "Do not translate correspondents or dates" in prompt def test_get_language_name_falls_back_to_language_code(): """ GIVEN: - A language code not present in settings.LANGUAGES WHEN: - get_language_name() is called THEN: - The original language code is returned unchanged """ assert get_language_name("zz-zz") == "zz-zz" def test_build_localization_prompt_preserves_unicode_characters(): """ GIVEN: - Suggestions containing non-ASCII characters WHEN: - build_localization_prompt() is called THEN: - The unicode characters are preserved as-is rather than escaped """ prompt = build_localization_prompt( { "title": "Gebührenbescheid", "tags": {"existing_ids": [], "new_names": []}, "correspondents": {"existing_ids": [], "new_names": []}, "document_types": {"existing_ids": [], "new_names": []}, "storage_paths": {"existing_ids": [], "new_names": []}, "dates": [], }, output_language="de-de", ) assert "Gebührenbescheid" in prompt assert "\\u00fc" not in prompt @pytest.mark.django_db def test_get_taxonomy_context_assembles_rag_text_and_candidates(): """ GIVEN: - A neighbour document with a tag, retrieved via retrieve_similar_nodes WHEN: - get_taxonomy_context() is called THEN: - The neighbour's tag appears in the taxonomy candidates - The neighbour's title/content appear in the RAG text context - The document's own (empty) assigned metadata is returned """ tag = TagFactory.create(name="Bloodwork") neighbour = DocumentFactory.create( content="Content of neighbour document", title="Neighbour Title", ) neighbour.tags.add(tag) document = DocumentFactory.create(content="Some content") fake_node = SimpleNamespace( metadata={"document_id": str(neighbour.pk)}, score=0.8, ) with patch( "paperless_ai.ai_classifier.retrieve_similar_nodes", return_value=[fake_node], ): candidates, assigned, context = get_taxonomy_context(document, user=None) assert candidates["tags"][0]["name"] == "Bloodwork" assert "TITLE: Neighbour Title" in context assert "Content of neighbour document" in context assert assigned == { "tags": [], "document_type": None, "correspondent": None, "storage_path": None, } @pytest.mark.django_db def test_get_taxonomy_context_no_similar_docs(): """ GIVEN: - No similar documents are retrieved WHEN: - get_taxonomy_context() is called THEN: - An empty RAG context and empty taxonomy candidates are returned """ document = DocumentFactory.create(content="Some content") with patch("paperless_ai.ai_classifier.retrieve_similar_nodes", return_value=[]): candidates, _assigned, context = get_taxonomy_context(document, user=None) assert context == "" assert candidates == { "tags": [], "document_types": [], "correspondents": [], "storage_paths": [], } 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 user at all) gets document_ids=None (no restriction) straight through to retrieve_similar_nodes(), instead of a full-library IN filter that is wasteful at best and, past ~32,763 documents, a hard sqlite3.OperationalError at worst (SQLite's bound-parameter limit). Ports the coverage that used to live on get_context_for_document before this refactor folded it into get_taxonomy_context. """ @pytest.mark.django_db def test_skips_permission_lookup_for_superuser( self, mocker: pytest_mock.MockerFixture, ) -> None: """ GIVEN: - A superuser WHEN: - get_taxonomy_context() is called THEN: - Permission lookup is skipped and no document_ids restriction is passed to retrieve_similar_nodes() """ document = DocumentFactory.create(content="Some content") mock_retrieve = mocker.patch( "paperless_ai.ai_classifier.retrieve_similar_nodes", return_value=[], ) mock_get_objects = mocker.patch( "paperless_ai.ai_classifier.get_objects_for_user_owner_aware", ) user = UserFactory.create(is_superuser=True) get_taxonomy_context(document, user) mock_get_objects.assert_not_called() assert mock_retrieve.call_args.kwargs["document_ids"] is None @pytest.mark.django_db def test_skips_permission_lookup_when_no_user( self, mocker: pytest_mock.MockerFixture, ) -> None: """ GIVEN: - No user is supplied WHEN: - get_taxonomy_context() is called THEN: - Permission lookup is skipped and no document_ids restriction is passed to retrieve_similar_nodes() """ document = DocumentFactory.create(content="Some content") mock_retrieve = mocker.patch( "paperless_ai.ai_classifier.retrieve_similar_nodes", return_value=[], ) mock_get_objects = mocker.patch( "paperless_ai.ai_classifier.get_objects_for_user_owner_aware", ) get_taxonomy_context(document, None) mock_get_objects.assert_not_called() assert mock_retrieve.call_args.kwargs["document_ids"] is None @pytest.mark.django_db def test_restricts_to_visible_documents_for_non_superuser( self, mocker: pytest_mock.MockerFixture, ) -> None: """ GIVEN: - A non-superuser WHEN: - get_taxonomy_context() is called THEN: - The user's visible document ids are looked up and passed to retrieve_similar_nodes() as a restriction """ document = DocumentFactory.create(content="Some content") mock_retrieve = mocker.patch( "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, ) user = UserFactory.create(is_superuser=False) get_taxonomy_context(document, user) mock_get_objects.assert_called_once_with(user, "view_document", Document) assert mock_retrieve.call_args.kwargs["document_ids"] == [1, 2, 3] @pytest.mark.django_db @patch("paperless_ai.ai_classifier.retrieve_similar_nodes") def test_get_taxonomy_context_retrieval_failure_degrades_to_no_hints(mock_retrieve): """ GIVEN: - retrieve_similar_nodes() raises an exception (e.g. vector store outage) WHEN: - get_taxonomy_context() is called THEN: - Empty taxonomy candidates and an empty RAG context are returned instead of propagating the exception """ document = DocumentFactory.create(content="Some content") mock_retrieve.side_effect = RuntimeError("vector store unavailable") candidates, _assigned, rag_context = get_taxonomy_context(document, user=None) assert candidates == { "tags": [], "document_types": [], "correspondents": [], "storage_paths": [], } assert rag_context == "" @pytest.mark.django_db @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( mock_retrieve, mock_build_candidates, ): """ GIVEN: - retrieve_similar_nodes() succeeds but build_taxonomy_candidates() raises (e.g. a DB or permission-backend failure) WHEN: - get_taxonomy_context() is called THEN: - Empty taxonomy candidates and an empty RAG context are returned instead of propagating the exception - the error boundary covers everything derived from the retrieval, not just the retrieval call itself """ document = DocumentFactory.create(content="Some content") mock_retrieve.return_value = [] mock_build_candidates.side_effect = RuntimeError("permission backend unavailable") candidates, _assigned, rag_context = get_taxonomy_context(document, user=None) assert candidates == { "tags": [], "document_types": [], "correspondents": [], "storage_paths": [], } assert rag_context == "" @pytest.mark.django_db def test_build_prompt_without_rag_includes_taxonomy_block(): """ GIVEN: - Non-empty taxonomy candidates WHEN: - build_prompt_without_rag() is called with candidates and assigned metadata THEN: - The candidate's id and the existing_ids instruction appear in the prompt """ document = DocumentFactory.create(content="Some content") config = AIConfig() candidates = { "tags": [{"id": 12, "name": "Bloodwork", "weight": 1.0}], "document_types": [], "correspondents": [], "storage_paths": [], } assigned = { "tags": [], "document_type": None, "correspondent": None, "storage_path": None, } prompt = build_prompt_without_rag( document, config, candidates=candidates, assigned=assigned, ) assert '"id": 12' in prompt assert "existing_ids" in prompt @pytest.mark.django_db def test_build_prompt_without_rag_identical_when_no_hints(): """ GIVEN: - Empty taxonomy candidates and empty assigned metadata WHEN: - build_prompt_without_rag() is called with those empty values, and separately with no candidates/assigned at all THEN: - Both prompts are identical - Neither mentions existing_ids or the "Available ..." candidate block: without any candidates in the prompt, that instruction would only invite the model to invent a plausible id that resolves to a real but unrelated object """ document = DocumentFactory.create(content="Some content") config = AIConfig() empty_candidates = { "tags": [], "document_types": [], "correspondents": [], "storage_paths": [], } empty_assigned = { "tags": [], "document_type": None, "correspondent": None, "storage_path": None, } with_empty_hints = build_prompt_without_rag( document, config, candidates=empty_candidates, assigned=empty_assigned, ) with_no_hints = build_prompt_without_rag(document, config) assert with_empty_hints == with_no_hints assert "existing_ids" not in with_no_hints assert "Available " not in with_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 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 THEN: - The localized new_names are used - The ORIGINAL existing_ids are kept, never the localized response's existing_ids - localization must never corrupt an exact taxonomy match """ 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 = [ { "title": "Invoice", "tags": {"existing_ids": [12], "new_names": ["Contractor Work"]}, "correspondents": {"existing_ids": [], "new_names": []}, "document_types": {"existing_ids": [], "new_names": []}, "storage_paths": {"existing_ids": [], "new_names": []}, "dates": [], }, { # The model's own localized-response existing_ids (999) must be # discarded - the merge always keeps the ORIGINAL resolved id. "title": "Rechnung", "tags": {"existing_ids": [999], "new_names": ["Auftragsarbeit"]}, "correspondents": {"existing_ids": [], "new_names": []}, "document_types": {"existing_ids": [], "new_names": []}, "storage_paths": {"existing_ids": [], "new_names": []}, "dates": [], }, ] result = get_ai_document_classification(document, output_language="de-de") localization_prompt = mock_client.run_llm_query.call_args_list[1].args[0] 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"]