From 9c475e0b27a0f1f272a0ed1dfa22ff5cff12de52 Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Sat, 29 Aug 2026 13:02:05 -0700 Subject: [PATCH] Fix: 3.1.0 llm suggestion raw cache user scoping (#13849) --- src/documents/caching.py | 66 +++++++++++++++--- src/documents/signals/handlers.py | 5 +- src/documents/tests/test_views.py | 112 +++++++++++++++++++++++++++--- src/documents/views.py | 9 ++- 4 files changed, 168 insertions(+), 24 deletions(-) diff --git a/src/documents/caching.py b/src/documents/caching.py index 162df92f5..e877d5827 100644 --- a/src/documents/caching.py +++ b/src/documents/caching.py @@ -1,7 +1,9 @@ from __future__ import annotations +import hashlib import logging import pickle +import uuid from binascii import hexlify from collections import OrderedDict from dataclasses import dataclass @@ -55,6 +57,8 @@ LLM_CACHE_CLASSIFIER_VERSION: Final[int] = 1001 CACHE_1_MINUTE: Final[int] = 60 CACHE_5_MINUTES: Final[int] = 5 * CACHE_1_MINUTE CACHE_50_MINUTES: Final[int] = 50 * CACHE_1_MINUTE +# Deliberately longer than any entry it names +LLM_CACHE_GENERATION_TIMEOUT: Final[int] = 2 * CACHE_50_MINUTES read_cache = caches["read-cache"] @@ -206,12 +210,40 @@ def refresh_suggestions_cache( cache.touch(doc_key, timeout) +def invalidate_suggestions_cache(document_id: int) -> None: + """Invalidate classifier-generated suggestions for a document.""" + cache.delete(get_suggestion_cache_key(document_id)) + + +def _llm_generation_key(document_id: int) -> str: + return f"{get_suggestion_cache_key(document_id)}_llm_generation" + + +def _llm_variant_key(document_id: int, backend: str) -> str: + """Cache key for one LLM configuration and permission scope. + + ``backend`` identifies the variant - model, endpoint, output language and + requesting user. + + Generating the token on first use lets invalidate_llm_suggestions_cache() + be no-op for documents that never had AI suggestions. + """ + generation_key = _llm_generation_key(document_id) + generation = cache.get_or_set( + generation_key, + lambda: uuid.uuid4().hex, + timeout=LLM_CACHE_GENERATION_TIMEOUT, + ) + cache.touch(generation_key, LLM_CACHE_GENERATION_TIMEOUT) + backend_hash = hashlib.sha256(backend.encode()).hexdigest()[:16] + return f"{get_suggestion_cache_key(document_id)}_llm_{generation}_{backend_hash}" + + def get_llm_suggestion_cache( document_id: int, backend: str, ) -> SuggestionCacheData | None: - doc_key = get_suggestion_cache_key(document_id) - data: SuggestionCacheData = cache.get(doc_key) + data: SuggestionCacheData = cache.get(_llm_variant_key(document_id, backend)) if ( data @@ -234,9 +266,8 @@ def set_llm_suggestions_cache( Cache LLM-generated suggestions using a backend-specific identifier (e.g. 'openai-like:gpt-4'). """ - doc_key = get_suggestion_cache_key(document_id) cache.set( - doc_key, + _llm_variant_key(document_id, backend), SuggestionCacheData( classifier_version=LLM_CACHE_CLASSIFIER_VERSION, classifier_hash=backend, @@ -246,17 +277,31 @@ def set_llm_suggestions_cache( ) +def refresh_llm_suggestions_cache( + document_id: int, + backend: str, + *, + timeout: int = CACHE_50_MINUTES, +) -> None: + """ + Refreshes the expiration of one cached LLM suggestion variant. + """ + cache.touch(_llm_variant_key(document_id, backend), timeout) + + def invalidate_llm_suggestions_cache( document_id: int, ) -> None: """ - Invalidate the LLM suggestions cache for a specific document and backend. + Invalidate every LLM suggestion variant for a document. """ - doc_key = get_suggestion_cache_key(document_id) - data: SuggestionCacheData = cache.get(doc_key) - - if data: - cache.delete(doc_key) + generation_key = _llm_generation_key(document_id) + if cache.get(generation_key) is not None: + cache.set( + generation_key, + uuid.uuid4().hex, + timeout=LLM_CACHE_GENERATION_TIMEOUT, + ) def get_metadata_cache_key(document_id: int) -> str: @@ -357,3 +402,4 @@ def clear_document_caches(document_id: int) -> None: get_thumbnail_modified_key(document_id), ], ) + invalidate_llm_suggestions_cache(document_id) diff --git a/src/documents/signals/handlers.py b/src/documents/signals/handlers.py index d6359066b..b76759781 100644 --- a/src/documents/signals/handlers.py +++ b/src/documents/signals/handlers.py @@ -32,6 +32,7 @@ from rest_framework import serializers from documents import matching from documents.caching import clear_document_caches from documents.caching import invalidate_llm_suggestions_cache +from documents.caching import invalidate_suggestions_cache from documents.data_models import ConsumableDocument from documents.file_handling import create_source_path_directory from documents.file_handling import delete_empty_directories @@ -740,9 +741,9 @@ def cleanup_custom_field_deletion(sender, instance: CustomField, **kwargs) -> No @receiver(models.signals.post_save, sender=Document) def update_llm_suggestions_cache(sender, instance, **kwargs): """ - Invalidate the LLM suggestions cache when a document is saved. + Invalidate suggestions caches when a document is saved. """ - # Invalidate the cache for the document + invalidate_suggestions_cache(instance.pk) invalidate_llm_suggestions_cache(instance.pk) diff --git a/src/documents/tests/test_views.py b/src/documents/tests/test_views.py index 81331f1ff..282203866 100644 --- a/src/documents/tests/test_views.py +++ b/src/documents/tests/test_views.py @@ -9,6 +9,7 @@ from django.conf import settings from django.contrib.auth.models import Group from django.contrib.auth.models import Permission from django.contrib.auth.models import User +from django.core.cache import cache from django.db import connection from django.test import TestCase from django.test import override_settings @@ -18,6 +19,7 @@ from guardian.shortcuts import assign_perm from rest_framework import status from documents.caching import get_llm_suggestion_cache +from documents.caching import get_suggestion_cache_key from documents.caching import set_llm_suggestions_cache from documents.models import Correspondent from documents.models import Document @@ -342,7 +344,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase): super().setUp() @patch("documents.views.get_llm_suggestion_cache") - @patch("documents.views.refresh_suggestions_cache") + @patch("documents.views.refresh_llm_suggestions_cache") @override_settings( AI_ENABLED=True, LLM_BACKEND="mock_backend", @@ -383,12 +385,15 @@ class TestAISuggestions(DirectoriesMixin, TestCase): self.assertEqual(response.json()["tags"], [self.tag1.pk]) mock_get_cache.assert_called_once_with( self.document.pk, - backend="mock_backend", + backend=f"mock_backend:user={self.user.pk}", + ) + mock_refresh_cache.assert_called_once_with( + self.document.pk, + backend=f"mock_backend:user={self.user.pk}", ) - mock_refresh_cache.assert_called_once_with(self.document.pk) @patch("documents.views.get_llm_suggestion_cache") - @patch("documents.views.refresh_suggestions_cache") + @patch("documents.views.refresh_llm_suggestions_cache") @override_settings( AI_ENABLED=True, LLM_BACKEND="mock_backend", @@ -524,7 +529,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase): self.assertEqual( get_llm_suggestion_cache( self.document.pk, - backend="mock_backend:de-de", + backend=f"mock_backend:de-de:user={self.user.pk}", ).suggestions["title"], "KI Title", ) @@ -563,7 +568,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase): self.assertEqual( get_llm_suggestion_cache( self.document.pk, - backend="mock_backend:fr-fr", + backend=f"mock_backend:fr-fr:user={self.user.pk}", ).suggestions["title"], "Titre IA", ) @@ -600,7 +605,79 @@ class TestAISuggestions(DirectoriesMixin, TestCase): self.assertIsNotNone( get_llm_suggestion_cache( self.document.pk, - backend="mock_backend:model-a:http://endpoint-a", + backend=(f"mock_backend:model-a:http://endpoint-a:user={self.user.pk}"), + ), + ) + + @patch("documents.views.get_ai_document_classification") + @override_settings( + AI_ENABLED=True, + LLM_BACKEND="mock_backend", + ) + def test_ai_suggestions_cache_variants_coexist_per_requesting_user( + self, + mock_get_ai_classification, + ) -> None: + """ + GIVEN: + - One user has populated the document's LLM suggestion cache + - A second user requests suggestions for the same document and + backend + WHEN: + - The second request is made + THEN: + - The first user's prompt-derived result is not reused + - The classification runs with the second user's visibility + context without evicting the first user's result + """ + second_user = User.objects.create_superuser(username="second_user") + empty_choices = { + "tags": {"existing_ids": [], "new_names": []}, + "correspondents": {"existing_ids": [], "new_names": []}, + "document_types": {"existing_ids": [], "new_names": []}, + "storage_paths": {"existing_ids": [], "new_names": []}, + "dates": [], + } + mock_get_ai_classification.side_effect = [ + {"title": "First user's result", **empty_choices}, + {"title": "Second user's result", **empty_choices}, + ] + + self.client.force_login(user=self.user) + first_response = self.client.get( + f"/api/documents/{self.document.pk}/ai_suggestions/", + ) + self.client.force_login(user=second_user) + second_response = self.client.get( + f"/api/documents/{self.document.pk}/ai_suggestions/", + ) + self.client.force_login(user=self.user) + first_cached_response = self.client.get( + f"/api/documents/{self.document.pk}/ai_suggestions/", + ) + + self.assertEqual(first_response.json()["title"], "First user's result") + self.assertEqual(second_response.json()["title"], "Second user's result") + self.assertEqual( + first_cached_response.json()["title"], + "First user's result", + ) + self.assertEqual(mock_get_ai_classification.call_count, 2) + mock_get_ai_classification.assert_called_with( + self.document, + second_user, + None, + ) + self.assertIsNotNone( + get_llm_suggestion_cache( + self.document.pk, + backend=f"mock_backend:user={second_user.pk}", + ), + ) + self.assertIsNotNone( + get_llm_suggestion_cache( + self.document.pk, + backend=f"mock_backend:user={self.user.pk}", ), ) @@ -786,8 +863,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase): self.assertEqual(response.json()["tags"], []) self.assertEqual(response.json()["suggested_tags"], []) - def test_invalidate_suggestions_cache(self) -> None: - self.client.force_login(user=self.user) + def test_document_save_invalidates_all_suggestion_caches(self) -> None: suggestions = { "title": "AI Title", "tags": ["tag1", "tag2"], @@ -796,11 +872,18 @@ class TestAISuggestions(DirectoriesMixin, TestCase): "storage_paths": ["path1"], "dates": ["2023-01-01"], } + standard_cache_key = get_suggestion_cache_key(self.document.pk) + cache.set(standard_cache_key, "classifier suggestions") set_llm_suggestions_cache( self.document.pk, suggestions, backend="mock_backend", ) + set_llm_suggestions_cache( + self.document.pk, + {**suggestions, "title": "Other Variant"}, + backend="other_backend:user=2", + ) self.assertEqual( get_llm_suggestion_cache( self.document.pk, @@ -808,17 +891,26 @@ class TestAISuggestions(DirectoriesMixin, TestCase): ).suggestions, suggestions, ) - # post_save signal triggered + self.assertEqual(cache.get(standard_cache_key), "classifier suggestions") + update_llm_suggestions_cache( sender=None, instance=self.document, ) + + self.assertIsNone(cache.get(standard_cache_key)) self.assertIsNone( get_llm_suggestion_cache( self.document.pk, backend="mock_backend", ), ) + self.assertIsNone( + get_llm_suggestion_cache( + self.document.pk, + backend="other_backend:user=2", + ), + ) class TestAIChatStreamingView(DirectoriesMixin, TestCase): diff --git a/src/documents/views.py b/src/documents/views.py index 09f5d914a..c6f0aeb15 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -113,6 +113,7 @@ from documents.bulk_download import OriginalsOnlyStrategy from documents.caching import get_llm_suggestion_cache from documents.caching import get_metadata_cache 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 @@ -1540,6 +1541,7 @@ class DocumentViewSet( ai_config.llm_model, ai_config.llm_endpoint, output_language, + f"user={request.user.pk}", ) if part ) @@ -1555,8 +1557,11 @@ class DocumentViewSet( # freshly for this requester on every request, cache hit or not, # so a resolved id cached for one user's visibility can never be # handed unfiltered to a second, less-privileged requester of - # the same (backend-keyed, not user-keyed) cache entry. - refresh_suggestions_cache(doc.pk) + # the same (backend + user-keyed) cache entry. + refresh_llm_suggestions_cache( + doc.pk, + backend=llm_cache_backend, + ) llm_suggestions = cached_llm_suggestions.suggestions else: try: