From aae947ede327497fbafc4d8828bb4cc5f09a00e2 Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Wed, 26 Aug 2026 08:24:43 -0700 Subject: [PATCH] Enhancement: prevent duplication of llm suggestion requests --- src/documents/caching.py | 60 ++++++++++++++++++++++ src/documents/tests/test_api_documents.py | 2 +- src/documents/tests/test_caching.py | 61 +++++++++++++++++++++++ src/documents/tests/test_views.py | 18 +++---- src/documents/views.py | 19 +++---- 5 files changed, 139 insertions(+), 21 deletions(-) diff --git a/src/documents/caching.py b/src/documents/caching.py index 559e83acb..d692ddfde 100644 --- a/src/documents/caching.py +++ b/src/documents/caching.py @@ -3,13 +3,16 @@ from __future__ import annotations import hashlib import logging import pickle +import time import uuid 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 @@ -21,6 +24,7 @@ from paperless.signed_pickle import signed_pickle_dumps from paperless.signed_pickle import signed_pickle_loads if TYPE_CHECKING: + from django.contrib.auth.models import User from django.core.cache.backends.base import BaseCache from documents.classifier import DocumentClassifier @@ -59,6 +63,9 @@ CLASSIFIER_MODIFIED_KEY: Final[str] = "classifier_modified" # validated separately, so candidate-anchored 1001 results are stale LLM_CACHE_CLASSIFIER_VERSION: Final[int] = 1002 +# 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 @@ -262,6 +269,59 @@ 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 + + lock_key = ( + f"{get_suggestion_cache_key(document.pk)}_llm_lock_" + f"{sha256(backend.encode()).hexdigest()}" + ) + + 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): + 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: + # Do not remove a replacement lock if this one expired while + # generation was still running. + if cache.get(lock_key) == lock_token: + cache.delete(lock_key) + + # 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, diff --git a/src/documents/tests/test_api_documents.py b/src/documents/tests/test_api_documents.py index 2aceb3fc9..9f7bed944 100644 --- a/src/documents/tests/test_api_documents.py +++ b/src/documents/tests/test_api_documents.py @@ -2608,7 +2608,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, diff --git a/src/documents/tests/test_caching.py b/src/documents/tests/test_caching.py index 41de11285..dccfc9a68 100644 --- a/src/documents/tests/test_caching.py +++ b/src/documents/tests/test_caching.py @@ -1,4 +1,9 @@ +from concurrent.futures import ThreadPoolExecutor +from threading import Event +from threading import Lock + from documents.caching import StoredLRUCache +from documents.caching import retrieve_llm_suggestions from paperless.signed_pickle import HMAC_SIZE from paperless.signed_pickle import signed_pickle_dumps from paperless.signed_pickle import signed_pickle_loads @@ -56,3 +61,59 @@ def test_stored_lru_cache_rejects_tampered_data(mocker) -> None: cache.load() assert cache.get("x") is None + + +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) diff --git a/src/documents/tests/test_views.py b/src/documents/tests/test_views.py index 282203866..2314c4256 100644 --- a/src/documents/tests/test_views.py +++ b/src/documents/tests/test_views.py @@ -446,7 +446,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", @@ -496,7 +496,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", @@ -534,7 +534,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", @@ -573,7 +573,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", @@ -681,7 +681,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", @@ -710,7 +710,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", @@ -737,7 +737,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", @@ -775,7 +775,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", @@ -814,7 +814,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", diff --git a/src/documents/views.py b/src/documents/views.py index a65c048e3..52419e81f 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -116,7 +116,7 @@ from documents.caching import get_suggestion_cache from documents.caching import refresh_llm_suggestions_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 @@ -249,7 +249,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 @@ -1575,10 +1574,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( @@ -1603,11 +1605,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"]