From 34a01a5bd69f62b3a690eb12d11662fdb5102144 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 | 55 ++++++++++++++++++++ src/documents/tests/test_api_documents.py | 2 +- src/documents/tests/test_caching.py | 63 +++++++++++++++++++++++ src/documents/tests/test_views.py | 18 +++---- src/documents/views.py | 13 ++--- 5 files changed, 133 insertions(+), 18 deletions(-) diff --git a/src/documents/caching.py b/src/documents/caching.py index 162df92f5..1be1159e8 100644 --- a/src/documents/caching.py +++ b/src/documents/caching.py @@ -2,20 +2,25 @@ 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 from django.core.cache import caches from documents.models import Document +from paperless_ai.ai_classifier import get_ai_document_classification if TYPE_CHECKING: + from django.contrib.auth.models import User from django.core.cache.backends.base import BaseCache from documents.classifier import DocumentClassifier @@ -223,6 +228,56 @@ 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.""" + 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) + 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) + 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(0.1) + + 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 b67b87772..520bd6a30 100644 --- a/src/documents/tests/test_api_documents.py +++ b/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("documents.caching.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 d75bda3c9..bb5eee52d 100644 --- a/src/documents/tests/test_caching.py +++ b/src/documents/tests/test_caching.py @@ -1,6 +1,12 @@ import pickle +from concurrent.futures import ThreadPoolExecutor +from threading import Event +from threading import Lock + +from django.core.cache import cache from documents.caching import StoredLRUCache +from documents.caching import retrieve_llm_suggestions def test_lru_cache_entries() -> None: @@ -43,3 +49,60 @@ 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: + cache.clear() + 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( + "documents.caching.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, + "ollama:model", + lock_timeout=10, + ) + assert generation_started.wait(timeout=2) + second = executor.submit( + retrieve_llm_suggestions, + document, + user, + None, + "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 81331f1ff..26d930f26 100644 --- a/src/documents/tests/test_views.py +++ b/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("documents.caching.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("documents.caching.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("documents.caching.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("documents.caching.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("documents.caching.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("documents.caching.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("documents.caching.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("documents.caching.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("documents.caching.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 bc77dc5eb..4024c38bb 100644 --- a/src/documents/views.py +++ b/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 @@ -246,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 @@ -1560,10 +1559,13 @@ class DocumentViewSet( llm_suggestions = cached_llm_suggestions.suggestions else: try: - llm_suggestions = get_ai_document_classification( + llm_suggestions = retrieve_llm_suggestions( doc, request.user, output_language, + llm_cache_backend, + # Classification, localization + 30s + lock_timeout=(2 * ai_config.llm_request_timeout) + 30, ) except ValueError as exc: logger.exception( @@ -1588,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"]