The LLM-generated test for that last thing

src/documents/caching.py

@@ -2,12 +2,15 @@ from __future__ import annotations
 
 import logging
 import pickle
+import time
 from binascii import hexlify
 from collections import OrderedDict
 from dataclasses import dataclass
+from hashlib import sha256
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Final
+from uuid import uuid4
 
 from django.conf import settings
 from django.core.cache import cache
@@ -16,6 +19,7 @@ from django.core.cache import caches
 from documents.models import Document
 
 if TYPE_CHECKING:
+    from django.contrib.auth.models import User
     from django.core.cache.backends.base import BaseCache
 
     from documents.classifier import DocumentClassifier
@@ -52,6 +56,9 @@ CLASSIFIER_MODIFIED_KEY: Final[str] = "classifier_modified"
 #          [...]} per taxonomy field (#13676)
 LLM_CACHE_CLASSIFIER_VERSION: Final[int] = 1001
 
+# How often a request waiting on llm generation re-checks the cache
+LLM_SUGGESTION_POLL_INTERVAL: Final[float] = 0.5
+
 CACHE_1_MINUTE: Final[int] = 60
 CACHE_5_MINUTES: Final[int] = 5 * CACHE_1_MINUTE
 CACHE_50_MINUTES: Final[int] = 50 * CACHE_1_MINUTE
@@ -223,6 +230,68 @@ def get_llm_suggestion_cache(
     return None
 
 
+def retrieve_llm_suggestions(
+    document: Document,
+    user: User | None,
+    output_language: str | None,
+    *,
+    backend: str,
+    lock_timeout: int,
+) -> dict:
+    """Return cached LLM suggestions, generating them once across workers."""
+    # Lazy import to avoid pulling in the whole AI stuff
+    from paperless_ai.ai_classifier import get_ai_document_classification
+    from paperless_ai.exceptions import LLMTimeoutError
+
+    lock_key = (
+        f"{get_suggestion_cache_key(document.pk)}_llm_lock_"
+        f"{sha256(backend.encode()).hexdigest()}"
+    )
+    waited = False
+
+    while True:
+        cached = get_llm_suggestion_cache(document.pk, backend=backend)
+        if cached is not None:
+            refresh_suggestions_cache(document.pk)
+            return cached.suggestions
+
+        lock_token = uuid4().hex
+        if cache.add(lock_key, lock_token, lock_timeout):
+            if waited:
+                # The generation we were waiting on has ended without caching
+                # anything so it either failed or outlived its lock. Give up
+                # rather than re-running it
+                cache.delete(lock_key)
+                raise LLMTimeoutError
+
+            try:
+                # The cache may have been populated while acquiring the lock.
+                cached = get_llm_suggestion_cache(document.pk, backend=backend)
+                if cached is not None:
+                    refresh_suggestions_cache(document.pk)
+                    return cached.suggestions
+
+                suggestions = get_ai_document_classification(
+                    document,
+                    user,
+                    output_language,
+                )
+                set_llm_suggestions_cache(
+                    document.pk,
+                    suggestions,
+                    backend=backend,
+                )
+                return suggestions
+            finally:
+                # Don't remove lock if this one expired while generation was still running
+                if cache.get(lock_key) == lock_token:
+                    cache.delete(lock_key)
+
+        waited = True
+        # Another worker is generating suggestions, poll to avoid another LLM request
+        time.sleep(LLM_SUGGESTION_POLL_INTERVAL)
+
+
 def set_llm_suggestions_cache(
     document_id: int,
     suggestions: dict,

src/documents/filters.py

@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING
 from typing import Any
 
 from django.contrib.contenttypes.models import ContentType
+from django.core.exceptions import FieldError
 from django.db.models import Case
 from django.db.models import CharField
 from django.db.models import Count
@@ -50,7 +51,6 @@ from documents.models import ShareLinkBundle
 from documents.models import StoragePath
 from documents.models import Tag
 from documents.permissions import permitted_object_ids
-from documents.versioning import ensure_effective_content
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -180,9 +180,14 @@ class TitleContentFilter(Filter):
             logger.warning(
                 "Deprecated document filter parameter 'title_content' used; use `text` instead.",
             )
-            return ensure_effective_content(qs).filter(
-                Q(title__icontains=value) | Q(effective_content__icontains=value),
-            )
+            try:
+                return qs.filter(
+                    Q(title__icontains=value) | Q(effective_content__icontains=value),
+                )
+            except FieldError:
+                return qs.filter(
+                    Q(title__icontains=value) | Q(content__icontains=value),
+                )
         else:
             return qs
 
@@ -193,9 +198,14 @@ class EffectiveContentFilter(Filter):
         value = value.strip() if isinstance(value, str) else value
         if not value:
             return qs
-        return ensure_effective_content(qs).filter(
-            **{f"effective_content__{self.lookup_expr}": value},
-        )
+        try:
+            return qs.filter(
+                **{f"effective_content__{self.lookup_expr}": value},
+            )
+        except FieldError:
+            return qs.filter(
+                **{f"content__{self.lookup_expr}": value},
+            )
 
 
 @extend_schema_field(serializers.BooleanField)

src/documents/tests/test_api_document_versions.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 import datetime
 from typing import TYPE_CHECKING
+from unittest import TestCase
 from unittest import mock
 
 from auditlog.models import LogEntry  # type: ignore[import-untyped]
 from django.contrib.auth.models import Permission
 from django.contrib.auth.models import User
 from django.contrib.contenttypes.models import ContentType
+from django.core.exceptions import FieldError
 from django.core.files.uploadedfile import SimpleUploadedFile
 from django.test import TestCase as DjangoTestCase
 from django.utils import timezone
@@ -20,7 +22,6 @@ from documents.filters import TitleContentFilter
 from documents.models import Document
 from documents.tests.utils import DirectoriesMixin
 from documents.tests.utils import read_streaming_response
-from documents.versioning import annotate_effective_content
 from documents.views import DocumentSelectionMixin
 
 if TYPE_CHECKING:
@@ -890,103 +891,32 @@ class TestDocumentVersioningApi(DirectoriesMixin, APITestCase):
         )
 
 
-class TestVersionAwareFilters(DjangoTestCase):
-    """
-    The filters annotate effective_content themselves rather than relying on
-    the caller's queryset carrying it, so they stay version-aware on a plain
-    Document queryset (e.g. the bulk-edit "select all matching" path).
-    """
+class TestVersionAwareFilters(TestCase):
+    def test_title_content_filter_falls_back_to_content(self) -> None:
+        queryset = mock.Mock()
+        fallback_queryset = mock.Mock()
+        queryset.filter.side_effect = [FieldError("missing field"), fallback_queryset]
 
-    def setUp(self) -> None:
-        super().setUp()
-        self.root = Document.objects.create(
-            title="root",
-            checksum="root",
-            mime_type="application/pdf",
-            content="superseded-content",
-        )
-        Document.objects.create(
-            title="version",
-            checksum="version",
-            mime_type="application/pdf",
-            root_document=self.root,
-            version_index=1,
-            content="latest-content",
-        )
-        self.unversioned = Document.objects.create(
-            title="unversioned",
-            checksum="unversioned",
-            mime_type="application/pdf",
-            content="latest-content",
-        )
+        result = TitleContentFilter().filter(queryset, " latest ")
 
-    def test_title_content_filter_matches_latest_version_content(self) -> None:
-        result = TitleContentFilter().filter(
-            Document.objects.filter(root_document__isnull=True),
+        self.assertIs(result, fallback_queryset)
+        self.assertEqual(queryset.filter.call_count, 2)
+
+    def test_effective_content_filter_falls_back_to_content_lookup(self) -> None:
+        queryset = mock.Mock()
+        fallback_queryset = mock.Mock()
+        queryset.filter.side_effect = [FieldError("missing field"), fallback_queryset]
+
+        result = EffectiveContentFilter(lookup_expr="icontains").filter(
+            queryset,
             " latest ",
         )
 
-        self.assertCountEqual(
-            [doc.id for doc in result],
-            [self.root.id, self.unversioned.id],
-        )
-
-    def test_effective_content_filter_matches_latest_version_content(self) -> None:
-        result = EffectiveContentFilter(lookup_expr="icontains").filter(
-            Document.objects.filter(root_document__isnull=True),
-            " latest ",
-        )
-
-        self.assertCountEqual(
-            [doc.id for doc in result],
-            [self.root.id, self.unversioned.id],
-        )
-
-    def test_effective_content_filter_ignores_superseded_content(self) -> None:
-        result = EffectiveContentFilter(lookup_expr="icontains").filter(
-            Document.objects.filter(root_document__isnull=True),
-            "superseded",
-        )
-
-        self.assertEqual(list(result), [])
-
-    def test_filters_reuse_an_existing_annotation(self) -> None:
-        """
-        Annotating twice under the same alias is an error, so an already
-        annotated queryset (the search path) has to be left alone.
-        """
-        annotated = annotate_effective_content(
-            Document.objects.filter(root_document__isnull=True),
-        )
-
-        result = EffectiveContentFilter(lookup_expr="icontains").filter(
-            annotated,
-            "latest",
-        )
-
-        self.assertCountEqual(
-            [doc.id for doc in result],
-            [self.root.id, self.unversioned.id],
-        )
-
-    def test_bulk_selection_does_not_match_superseded_content(self) -> None:
-        """
-        Bulk edit's "select all matching" builds its own queryset, so before
-        the filters annotated for themselves it matched the root document's
-        superseded content -- selecting documents the list view, filtered by
-        the same term, does not show.
-        """
-        user = User.objects.create_superuser(username="bulk_selection")
-
-        selected = DocumentSelectionMixin()._resolve_document_ids(
-            user=user,
-            validated_data={
-                "all": True,
-                "filters": {"content__icontains": "superseded"},
-            },
-        )
-
-        self.assertEqual(selected, [])
+        self.assertIs(result, fallback_queryset)
+        first_kwargs = queryset.filter.call_args_list[0].kwargs
+        second_kwargs = queryset.filter.call_args_list[1].kwargs
+        self.assertEqual(first_kwargs, {"effective_content__icontains": "latest"})
+        self.assertEqual(second_kwargs, {"content__icontains": "latest"})
 
     def test_effective_content_filter_returns_input_for_empty_values(self) -> None:
         queryset = mock.Mock()

src/documents/tests/test_api_documents.py

@@ -2486,7 +2486,7 @@ class TestDocumentApi(DirectoriesMixin, ConsumeTaskMixin, APITestCase):
         response = self.client.get("/api/documents/34676/suggestions/")
         self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
 
-    @mock.patch("documents.views.get_ai_document_classification")
+    @mock.patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(AI_ENABLED=True)
     def test_suggestions_still_uses_classifier_when_ai_enabled(
         self,

src/documents/tests/test_api_search.py

@@ -1917,29 +1917,6 @@ class TestDocumentSearchApi(DirectoriesMixin, APITestCase):
         self.assertEqual(len(response.data["documents"]), 1)
         self.assertEqual(response.data["documents"][0]["id"], title_match.id)
 
-    def test_global_search_returns_latest_version_content(self) -> None:
-        root = Document.objects.create(
-            title="bank statement",
-            content="superseded content",
-            checksum="GSV1",
-            pk=23,
-        )
-        Document.objects.create(
-            title="bank statement v2",
-            content="latest content",
-            checksum="GSV2",
-            pk=24,
-            root_document=root,
-            version_index=1,
-        )
-
-        self.client.force_authenticate(self.user)
-
-        response = self.client.get("/api/search/?query=bank&db_only=true")
-        self.assertEqual(response.status_code, status.HTTP_200_OK)
-        returned = {doc["id"]: doc["content"] for doc in response.data["documents"]}
-        self.assertEqual(returned.get(root.id), "latest content")
-
     def test_global_search_filters_owned_mail_objects(self) -> None:
         user1 = User.objects.create_user("mail-search-user")
         user2 = User.objects.create_user("other-mail-search-user")

src/documents/tests/test_caching.py

@@ -1,6 +1,13 @@
 import pickle
+from concurrent.futures import ThreadPoolExecutor
+from threading import Event
+from threading import Lock
+
+import pytest
 
 from documents.caching import StoredLRUCache
+from documents.caching import retrieve_llm_suggestions
+from paperless_ai.exceptions import LLMTimeoutError
 
 
 def test_lru_cache_entries() -> None:
@@ -43,3 +50,119 @@ def test_stored_lru_cache_key_ttl(mocker) -> None:
     assert key == "test_key"
     assert timeout == 321
     assert pickle.loads(data) == {"x": "X", "y": "Y"}
+
+
+def test_llm_suggestions_are_generated_once_for_concurrent_requests(mocker) -> None:
+    generation_started = Event()
+    finish_generation = Event()
+    waiter_started = Event()
+    call_lock = Lock()
+    calls = 0
+    suggestions = {"title": "Generated once"}
+    document = mocker.Mock(pk=42)
+    user = mocker.Mock()
+
+    def generate(*args) -> dict:
+        nonlocal calls
+        with call_lock:
+            calls += 1
+        generation_started.set()
+        assert finish_generation.wait(timeout=2)
+        return suggestions
+
+    def wait_for_generation(_interval: float) -> None:
+        waiter_started.set()
+        assert finish_generation.wait(timeout=2)
+
+    mock_get_classification = mocker.patch(
+        "paperless_ai.ai_classifier.get_ai_document_classification",
+        side_effect=generate,
+    )
+    mocker.patch("documents.caching.time.sleep", side_effect=wait_for_generation)
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        first = executor.submit(
+            retrieve_llm_suggestions,
+            document,
+            user,
+            None,
+            backend="ollama:model",
+            lock_timeout=10,
+        )
+        assert generation_started.wait(timeout=2)
+        second = executor.submit(
+            retrieve_llm_suggestions,
+            document,
+            user,
+            None,
+            backend="ollama:model",
+            lock_timeout=10,
+        )
+        assert waiter_started.wait(timeout=2)
+        finish_generation.set()
+
+        assert first.result(timeout=2) == suggestions
+        assert second.result(timeout=2) == suggestions
+
+    assert calls == 1
+    mock_get_classification.assert_called_once_with(document, user, None)
+
+
+def test_llm_suggestions_waiter_does_not_rerun_a_failed_generation(mocker) -> None:
+    """
+    A request queued behind a generation that fails should give up, not take
+    its turn at re-running a query that just failed.
+    """
+    generation_started = Event()
+    fail_generation = Event()
+    waiter_started = Event()
+    call_lock = Lock()
+    calls = 0
+    document = mocker.Mock(pk=43)
+    user = mocker.Mock()
+
+    def generate(*args) -> dict:
+        nonlocal calls
+        with call_lock:
+            calls += 1
+        generation_started.set()
+        assert fail_generation.wait(timeout=2)
+        raise ValueError("Unknown model")
+
+    def wait_for_generation(_interval: float) -> None:
+        waiter_started.set()
+        assert fail_generation.wait(timeout=2)
+
+    mocker.patch(
+        "paperless_ai.ai_classifier.get_ai_document_classification",
+        side_effect=generate,
+    )
+    mocker.patch("documents.caching.time.sleep", side_effect=wait_for_generation)
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        first = executor.submit(
+            retrieve_llm_suggestions,
+            document,
+            user,
+            None,
+            backend="ollama:model",
+            lock_timeout=10,
+        )
+        assert generation_started.wait(timeout=2)
+        second = executor.submit(
+            retrieve_llm_suggestions,
+            document,
+            user,
+            None,
+            backend="ollama:model",
+            lock_timeout=10,
+        )
+        assert waiter_started.wait(timeout=2)
+        fail_generation.set()
+
+        with pytest.raises(ValueError, match="Unknown model"):
+            first.result(timeout=2)
+        with pytest.raises(LLMTimeoutError):
+            second.result(timeout=2)
+
+    assert calls == 1

src/documents/tests/test_views.py

@@ -441,7 +441,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
         self.assertEqual(response.json()["tags"], [])
         self.assertEqual(response.json()["suggested_tags"], [])
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -491,7 +491,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             None,
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -529,7 +529,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             "KI Title",
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -568,7 +568,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             "Titre IA",
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -604,7 +604,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             ),
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="openai-like",
@@ -633,7 +633,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             get_llm_suggestion_cache(self.document.pk, backend="openai-like"),
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="openai-like",
@@ -660,7 +660,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             get_llm_suggestion_cache(self.document.pk, backend="openai-like"),
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -698,7 +698,7 @@ 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")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -737,7 +737,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
         self.assertEqual(response.json()["tags"], [self.tag1.pk])
         self.assertEqual(response.json()["suggested_tags"], [])
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",

src/documents/versioning.py

@@ -43,21 +43,6 @@ def annotate_effective_content(documents: QuerySet[Document]) -> QuerySet[Docume
     )
 
 
-def ensure_effective_content(documents: QuerySet[Document]) -> QuerySet[Document]:
-    """
-    Annotates effective_content unless the queryset already carries it.
-
-    Lets a filter depend on effective_content without having to assume its
-    caller annotated one -- annotating twice under the same alias is an error,
-    and silently matching on the root document's own content instead is worse,
-    because the same filter then selects different documents depending on which
-    queryset it was handed.
-    """
-    if "effective_content" in documents.query.annotations:
-        return documents
-    return annotate_effective_content(documents)
-
-
 def sort_versions_newest_first(documents: list[Document]) -> list[Document]:
     """
     Same sorting as versions_newest_first()

src/documents/views.py

@@ -115,7 +115,7 @@ from documents.caching import get_metadata_cache
 from documents.caching import get_suggestion_cache
 from documents.caching import refresh_metadata_cache
 from documents.caching import refresh_suggestions_cache
-from documents.caching import set_llm_suggestions_cache
+from documents.caching import retrieve_llm_suggestions
 from documents.caching import set_metadata_cache
 from documents.caching import set_suggestions_cache
 from documents.classifier import load_classifier
@@ -230,7 +230,6 @@ from documents.tasks import train_classifier
 from documents.tasks import update_document_parent_tags
 from documents.utils import get_boolean
 from documents.versioning import VersionResolutionError
-from documents.versioning import annotate_effective_content
 from documents.versioning import get_latest_version_for_root
 from documents.versioning import get_request_version_param
 from documents.versioning import get_root_document
@@ -247,7 +246,6 @@ from paperless.parsers.remote import RemoteEngineConfig
 from paperless.serialisers import GroupSerializer
 from paperless.serialisers import UserSerializer
 from paperless.views import StandardPagination
-from paperless_ai.ai_classifier import get_ai_document_classification
 from paperless_ai.ai_classifier import get_llm_output_language
 from paperless_ai.chat import stream_chat_with_documents
 from paperless_ai.exceptions import LLMTimeoutError
@@ -1561,10 +1559,13 @@ class DocumentViewSet(
             llm_suggestions = cached_llm_suggestions.suggestions
         else:
             try:
-                llm_suggestions = get_ai_document_classification(
-                    doc,
-                    request.user,
-                    output_language,
+                llm_suggestions = retrieve_llm_suggestions(
+                    document=doc,
+                    user=request.user,
+                    output_language=output_language,
+                    backend=llm_cache_backend,
+                    # Classification, localization + 30s
+                    lock_timeout=(2 * ai_config.llm_request_timeout) + 30,
                 )
             except ValueError as exc:
                 logger.exception(
@@ -1589,11 +1590,6 @@ class DocumentViewSet(
                     {"ai": [_("AI backend request timed out.")]},
                     status=status.HTTP_503_SERVICE_UNAVAILABLE,
                 )
-            set_llm_suggestions_cache(
-                doc.pk,
-                llm_suggestions,
-                backend=llm_cache_backend,
-            )
 
         tags_choice: TaxonomyChoiceDict = llm_suggestions["tags"]
         correspondents_choice: TaxonomyChoiceDict = llm_suggestions["correspondents"]
@@ -3614,13 +3610,8 @@ class GlobalSearchView(PassUserMixin):
         OBJECT_LIMIT = 3
         docs = []
         if request.user.has_perm("documents.view_document"):
-            # Never more than OBJECT_LIMIT rows come back here, so annotating
-            # is cheap -- and without it these results show the root
-            # document's superseded content.
-            all_docs = annotate_effective_content(
-                Document.objects.filter(
-                    id__in=permitted_document_ids(request.user),
-                ),
+            all_docs = Document.objects.filter(
+                id__in=permitted_document_ids(request.user),
             )
             if db_only:
                 docs = all_docs.filter(title__icontains=query)[:OBJECT_LIMIT]