mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-08-29 14:07:33 +00:00
tidy up
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user