From 15a5618a6e3da23a9bcea3cd7d1a218080dfda1d Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Wed, 26 Aug 2026 09:19:27 -0700 Subject: [PATCH] tidy up --- src/documents/caching.py | 15 ++++++++++----- src/documents/tests/test_api_documents.py | 2 +- src/documents/tests/test_caching.py | 9 +++------ src/documents/tests/test_views.py | 18 +++++++++--------- src/documents/views.py | 8 ++++---- 5 files changed, 27 insertions(+), 25 deletions(-) diff --git a/src/documents/caching.py b/src/documents/caching.py index 1be1159e8..a0086b312 100644 --- a/src/documents/caching.py +++ b/src/documents/caching.py @@ -17,7 +17,6 @@ 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 @@ -57,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 @@ -232,18 +234,21 @@ def retrieve_llm_suggestions( document: Document, user: User | None, output_language: str | None, - backend: str, *, + 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) + cached = get_llm_suggestion_cache(document.pk, backend=backend) if cached is not None: refresh_suggestions_cache(document.pk) return cached.suggestions @@ -252,7 +257,7 @@ def retrieve_llm_suggestions( 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) + cached = get_llm_suggestion_cache(document.pk, backend=backend) if cached is not None: refresh_suggestions_cache(document.pk) return cached.suggestions @@ -275,7 +280,7 @@ def retrieve_llm_suggestions( cache.delete(lock_key) # Another worker is generating suggestions, poll to avoid another LLM request - time.sleep(0.1) + time.sleep(LLM_SUGGESTION_POLL_INTERVAL) def set_llm_suggestions_cache( diff --git a/src/documents/tests/test_api_documents.py b/src/documents/tests/test_api_documents.py index 520bd6a30..594ed8190 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.caching.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 bb5eee52d..32460a47c 100644 --- a/src/documents/tests/test_caching.py +++ b/src/documents/tests/test_caching.py @@ -3,8 +3,6 @@ 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 @@ -52,7 +50,6 @@ def test_stored_lru_cache_key_ttl(mocker) -> None: def test_llm_suggestions_are_generated_once_for_concurrent_requests(mocker) -> None: - cache.clear() generation_started = Event() finish_generation = Event() waiter_started = Event() @@ -75,7 +72,7 @@ def test_llm_suggestions_are_generated_once_for_concurrent_requests(mocker) -> N assert finish_generation.wait(timeout=2) mock_get_classification = mocker.patch( - "documents.caching.get_ai_document_classification", + "paperless_ai.ai_classifier.get_ai_document_classification", side_effect=generate, ) mocker.patch("documents.caching.time.sleep", side_effect=wait_for_generation) @@ -86,7 +83,7 @@ def test_llm_suggestions_are_generated_once_for_concurrent_requests(mocker) -> N document, user, None, - "ollama:model", + backend="ollama:model", lock_timeout=10, ) assert generation_started.wait(timeout=2) @@ -95,7 +92,7 @@ def test_llm_suggestions_are_generated_once_for_concurrent_requests(mocker) -> N document, user, None, - "ollama:model", + backend="ollama:model", lock_timeout=10, ) assert waiter_started.wait(timeout=2) diff --git a/src/documents/tests/test_views.py b/src/documents/tests/test_views.py index 26d930f26..10f9d409d 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.caching.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.caching.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.caching.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.caching.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.caching.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.caching.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.caching.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.caching.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.caching.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 4024c38bb..f1526223d 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -1560,10 +1560,10 @@ class DocumentViewSet( else: try: llm_suggestions = retrieve_llm_suggestions( - doc, - request.user, - output_language, - llm_cache_backend, + 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, )