mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-08-14 23:03:22 +00:00
AI Suggestions previously invented near-duplicate metadata because the classification prompt had no knowledge of the installation's own taxonomy. This surfaces a small, ranked, permission-filtered set of existing tags/document types/correspondents/storage paths - drawn from the document's RAG neighbors plus its own already-assigned metadata - so the model prefers reusing what already exists. The LLM response schema now returns existing_ids (IDs of reused candidates) separately from new_names (genuinely new suggestions). Only new_names goes through localization and fuzzy name-matching; existing_ids is resolved deterministically and never touched by the localization pass, so exact matches can no longer be silently corrupted by translation.
963 lines
35 KiB
Python
963 lines
35 KiB
Python
import json
|
|
import tempfile
|
|
from datetime import timedelta
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
from unittest.mock import patch
|
|
|
|
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.db import connection
|
|
from django.test import TestCase
|
|
from django.test import override_settings
|
|
from django.test.utils import CaptureQueriesContext
|
|
from django.utils import timezone
|
|
from guardian.shortcuts import assign_perm
|
|
from rest_framework import status
|
|
|
|
from documents.caching import get_llm_suggestion_cache
|
|
from documents.caching import set_llm_suggestions_cache
|
|
from documents.models import Correspondent
|
|
from documents.models import Document
|
|
from documents.models import DocumentType
|
|
from documents.models import ShareLink
|
|
from documents.models import StoragePath
|
|
from documents.models import Tag
|
|
from documents.models import UiSettings
|
|
from documents.signals.handlers import update_llm_suggestions_cache
|
|
from documents.tests.utils import DirectoriesMixin
|
|
from documents.tests.utils import read_streaming_response
|
|
from paperless.models import ApplicationConfiguration
|
|
from paperless_ai.exceptions import LLMTimeoutError
|
|
|
|
|
|
class TestViews(DirectoriesMixin, TestCase):
|
|
@classmethod
|
|
def setUpTestData(cls) -> None:
|
|
super().setUpTestData()
|
|
ApplicationConfiguration.objects.get_or_create()
|
|
|
|
def setUp(self) -> None:
|
|
self.user = User.objects.create_user("testuser")
|
|
super().setUp()
|
|
|
|
def test_login_redirect(self) -> None:
|
|
response = self.client.get("/")
|
|
self.assertEqual(response.status_code, status.HTTP_302_FOUND)
|
|
self.assertEqual(response.url, "/accounts/login/?next=/")
|
|
|
|
def test_index(self) -> None:
|
|
self.client.force_login(self.user)
|
|
for language_given, language_actual in [
|
|
("", "en-US"),
|
|
("en-US", "en-US"),
|
|
("de", "de-DE"),
|
|
("en", "en-US"),
|
|
("en-us", "en-US"),
|
|
("fr", "fr-FR"),
|
|
("jp", "en-US"),
|
|
]:
|
|
if language_given:
|
|
self.client.cookies.load(
|
|
{settings.LANGUAGE_COOKIE_NAME: language_given},
|
|
)
|
|
elif settings.LANGUAGE_COOKIE_NAME in self.client.cookies:
|
|
self.client.cookies.pop(settings.LANGUAGE_COOKIE_NAME)
|
|
|
|
response = self.client.get(
|
|
"/",
|
|
)
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(
|
|
response.context_data["webmanifest"],
|
|
f"frontend/{language_actual}/manifest.webmanifest",
|
|
)
|
|
self.assertEqual(
|
|
response.context_data["styles_css"],
|
|
f"frontend/{language_actual}/styles.css",
|
|
)
|
|
self.assertEqual(
|
|
response.context_data["polyfills_js"],
|
|
f"frontend/{language_actual}/polyfills.js",
|
|
)
|
|
self.assertEqual(
|
|
response.context_data["main_js"],
|
|
f"frontend/{language_actual}/main.js",
|
|
)
|
|
|
|
@override_settings(BASE_URL="/paperless/")
|
|
def test_index_app_logo_with_base_url(self) -> None:
|
|
"""
|
|
GIVEN:
|
|
- Existing config with app_logo specified
|
|
WHEN:
|
|
- Index page is loaded
|
|
THEN:
|
|
- app_logo is prefixed with BASE_URL
|
|
"""
|
|
config = ApplicationConfiguration.objects.first()
|
|
config.app_logo = "/logo/example.jpg"
|
|
config.save()
|
|
self.client.force_login(self.user)
|
|
response = self.client.get("/")
|
|
self.assertEqual(
|
|
response.context["APP_LOGO"],
|
|
f"/paperless{config.app_logo}",
|
|
)
|
|
|
|
def test_share_link_views(self) -> None:
|
|
"""
|
|
GIVEN:
|
|
- Share link created
|
|
WHEN:
|
|
- Valid request for share link is made
|
|
- Invalid request for share link is made
|
|
- Request for expired share link is made
|
|
THEN:
|
|
- Document is returned without need for login
|
|
- User is redirected to login with error
|
|
- User is redirected to login with error
|
|
"""
|
|
|
|
_, filename = tempfile.mkstemp(dir=self.dirs.originals_dir)
|
|
|
|
content = b"This is a test"
|
|
|
|
with Path(filename).open("wb") as f:
|
|
f.write(content)
|
|
|
|
doc = Document.objects.create(
|
|
title="none",
|
|
filename=Path(filename).name,
|
|
mime_type="application/pdf",
|
|
)
|
|
|
|
sharelink_permissions = Permission.objects.filter(
|
|
codename__contains="sharelink",
|
|
)
|
|
self.user.user_permissions.add(*sharelink_permissions)
|
|
self.user.save()
|
|
|
|
self.client.force_login(self.user)
|
|
|
|
self.client.post(
|
|
"/api/share_links/",
|
|
{
|
|
"document": doc.pk,
|
|
"file_version": "original",
|
|
},
|
|
)
|
|
sl1 = ShareLink.objects.get(document=doc)
|
|
|
|
self.client.logout()
|
|
|
|
# Valid
|
|
response = self.client.get(f"/share/{sl1.slug}")
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(read_streaming_response(response), content)
|
|
|
|
# Invalid
|
|
response = self.client.get("/share/123notaslug", follow=True)
|
|
response.render()
|
|
self.assertEqual(response.request["PATH_INFO"], "/accounts/login/")
|
|
self.assertContains(response, b"Share link was not found")
|
|
|
|
# Expired
|
|
sl1.expiration = timezone.now() - timedelta(days=1)
|
|
sl1.save()
|
|
|
|
response = self.client.get(f"/share/{sl1.slug}", follow=True)
|
|
response.render()
|
|
self.assertEqual(response.request["PATH_INFO"], "/accounts/login/")
|
|
self.assertContains(response, b"Share link has expired")
|
|
|
|
def test_share_link_archive_falls_back_to_original(self) -> None:
|
|
"""
|
|
GIVEN:
|
|
- A document without an archive version
|
|
- A share link using the default archive file version
|
|
WHEN:
|
|
- An unauthenticated request for the share link is made
|
|
THEN:
|
|
- The original document is returned
|
|
"""
|
|
_, filename = tempfile.mkstemp(dir=self.dirs.originals_dir)
|
|
content = b"This document has no archive"
|
|
|
|
with Path(filename).open("wb") as f:
|
|
f.write(content)
|
|
|
|
doc = Document.objects.create(
|
|
title="no archive",
|
|
filename=Path(filename).name,
|
|
mime_type="text/plain",
|
|
)
|
|
|
|
sharelink_permissions = Permission.objects.filter(
|
|
codename__contains="sharelink",
|
|
)
|
|
self.user.user_permissions.add(*sharelink_permissions)
|
|
self.client.force_login(self.user)
|
|
|
|
create_response = self.client.post(
|
|
"/api/share_links/",
|
|
{"document": doc.pk},
|
|
)
|
|
self.assertEqual(create_response.status_code, status.HTTP_201_CREATED)
|
|
share_link = ShareLink.objects.get(document=doc)
|
|
self.assertEqual(share_link.file_version, ShareLink.FileVersion.ARCHIVE)
|
|
|
|
self.client.logout()
|
|
|
|
response = self.client.get(f"/share/{share_link.slug}")
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(read_streaming_response(response), content)
|
|
|
|
def test_list_with_full_permissions(self) -> None:
|
|
"""
|
|
GIVEN:
|
|
- Tags with different permissions
|
|
WHEN:
|
|
- Request to get tag list with full permissions is made
|
|
THEN:
|
|
- Tag list is returned with the right permission information
|
|
"""
|
|
user2 = User.objects.create(username="user2")
|
|
user3 = User.objects.create(username="user3")
|
|
group1 = Group.objects.create(name="group1")
|
|
group2 = Group.objects.create(name="group2")
|
|
group3 = Group.objects.create(name="group3")
|
|
t1 = Tag.objects.create(name="invoice", pk=1)
|
|
assign_perm("view_tag", self.user, t1)
|
|
assign_perm("view_tag", user2, t1)
|
|
assign_perm("view_tag", user3, t1)
|
|
assign_perm("view_tag", group1, t1)
|
|
assign_perm("view_tag", group2, t1)
|
|
assign_perm("view_tag", group3, t1)
|
|
assign_perm("change_tag", self.user, t1)
|
|
assign_perm("change_tag", user2, t1)
|
|
assign_perm("change_tag", group1, t1)
|
|
assign_perm("change_tag", group2, t1)
|
|
|
|
Tag.objects.create(name="bank statement", pk=2)
|
|
d1 = Document.objects.create(
|
|
title="Invoice 1",
|
|
content="This is the invoice of a very expensive item",
|
|
checksum="A",
|
|
)
|
|
d1.tags.add(t1)
|
|
d2 = Document.objects.create(
|
|
title="Invoice 2",
|
|
content="Internet invoice, I should pay it to continue contributing",
|
|
checksum="B",
|
|
)
|
|
d2.tags.add(t1)
|
|
|
|
view_permissions = Permission.objects.filter(
|
|
codename__contains="view_tag",
|
|
)
|
|
self.user.user_permissions.add(*view_permissions)
|
|
self.user.save()
|
|
|
|
self.client.force_login(self.user)
|
|
response = self.client.get("/api/tags/?page=1&full_perms=true")
|
|
results = json.loads(response.content)["results"]
|
|
for tag in results:
|
|
if tag["name"] == "invoice":
|
|
assert tag["permissions"] == {
|
|
"view": {
|
|
"users": [self.user.pk, user2.pk, user3.pk],
|
|
"groups": [group1.pk, group2.pk, group3.pk],
|
|
},
|
|
"change": {
|
|
"users": [self.user.pk, user2.pk],
|
|
"groups": [group1.pk, group2.pk],
|
|
},
|
|
}
|
|
elif tag["name"] == "bank statement":
|
|
assert tag["permissions"] == {
|
|
"view": {"users": [], "groups": []},
|
|
"change": {"users": [], "groups": []},
|
|
}
|
|
else:
|
|
assert False, f"Unexpected tag found: {tag['name']}"
|
|
|
|
def test_list_no_n_plus_1_queries(self) -> None:
|
|
"""
|
|
GIVEN:
|
|
- Tags with different permissions
|
|
WHEN:
|
|
- Request to get tag list with full permissions is made
|
|
THEN:
|
|
- Permissions are not queried in database tag by tag,
|
|
i.e. there are no N+1 queries
|
|
"""
|
|
view_permissions = Permission.objects.filter(
|
|
codename__contains="view_tag",
|
|
)
|
|
self.user.user_permissions.add(*view_permissions)
|
|
self.user.save()
|
|
self.client.force_login(self.user)
|
|
|
|
# Start by a small list, and count the number of SQL queries
|
|
for i in range(2):
|
|
Tag.objects.create(name=f"tag_{i}")
|
|
|
|
with CaptureQueriesContext(connection) as ctx_small:
|
|
response_small = self.client.get("/api/tags/?full_perms=true")
|
|
assert response_small.status_code == 200
|
|
num_queries_small = len(ctx_small.captured_queries)
|
|
|
|
# Complete the list, and count the number of SQL queries again
|
|
for i in range(2, 50):
|
|
Tag.objects.create(name=f"tag_{i}")
|
|
|
|
with CaptureQueriesContext(connection) as ctx_large:
|
|
response_large = self.client.get("/api/tags/?full_perms=true")
|
|
assert response_large.status_code == 200
|
|
num_queries_large = len(ctx_large.captured_queries)
|
|
|
|
# A few additional queries are allowed, but not a linear explosion
|
|
assert num_queries_large <= num_queries_small + 5, (
|
|
f"Possible N+1 queries detected: {num_queries_small} queries for 2 tags, "
|
|
f"but {num_queries_large} queries for 50 tags"
|
|
)
|
|
|
|
|
|
class TestAISuggestions(DirectoriesMixin, TestCase):
|
|
def setUp(self) -> None:
|
|
self.user = User.objects.create_superuser(username="testuser")
|
|
self.document = Document.objects.create(
|
|
title="Test Document",
|
|
filename="test.pdf",
|
|
mime_type="application/pdf",
|
|
)
|
|
self.tag1 = Tag.objects.create(name="tag1")
|
|
self.correspondent1 = Correspondent.objects.create(name="correspondent1")
|
|
self.document_type1 = DocumentType.objects.create(name="type1")
|
|
self.path1 = StoragePath.objects.create(name="path1")
|
|
super().setUp()
|
|
|
|
@patch("documents.views.get_llm_suggestion_cache")
|
|
@patch("documents.views.refresh_suggestions_cache")
|
|
@override_settings(
|
|
AI_ENABLED=True,
|
|
LLM_BACKEND="mock_backend",
|
|
)
|
|
def test_ai_suggestions_with_cached_llm(
|
|
self,
|
|
mock_refresh_cache,
|
|
mock_get_cache,
|
|
) -> None:
|
|
"""
|
|
GIVEN:
|
|
- A cached LLM classification holding the raw existing_ids/
|
|
new_names choices (never resolved object ids)
|
|
WHEN:
|
|
- ai_suggestions is requested
|
|
THEN:
|
|
- The cached choices are resolved into ids for this request
|
|
(not returned verbatim from the cache) and the cache's TTL is
|
|
refreshed
|
|
"""
|
|
mock_get_cache.return_value = MagicMock(
|
|
suggestions={
|
|
"title": "Cached Title",
|
|
"tags": {"existing_ids": [self.tag1.pk], "new_names": []},
|
|
"correspondents": {"existing_ids": [], "new_names": []},
|
|
"document_types": {"existing_ids": [], "new_names": []},
|
|
"storage_paths": {"existing_ids": [], "new_names": []},
|
|
"dates": [],
|
|
},
|
|
)
|
|
|
|
self.client.force_login(user=self.user)
|
|
response = self.client.get(
|
|
f"/api/documents/{self.document.pk}/ai_suggestions/",
|
|
)
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(response.json()["title"], "Cached Title")
|
|
self.assertEqual(response.json()["tags"], [self.tag1.pk])
|
|
mock_get_cache.assert_called_once_with(
|
|
self.document.pk,
|
|
backend="mock_backend",
|
|
)
|
|
mock_refresh_cache.assert_called_once_with(self.document.pk)
|
|
|
|
@patch("documents.views.get_llm_suggestion_cache")
|
|
@patch("documents.views.refresh_suggestions_cache")
|
|
@override_settings(
|
|
AI_ENABLED=True,
|
|
LLM_BACKEND="mock_backend",
|
|
)
|
|
def test_ai_suggestions_cache_hit_re_filters_for_narrower_requester(
|
|
self,
|
|
mock_refresh_cache,
|
|
mock_get_cache,
|
|
) -> None:
|
|
"""
|
|
GIVEN:
|
|
- A cached LLM classification whose existing_ids include a tag
|
|
only visible to a broader-visibility user (e.g. the requester
|
|
who originally generated it)
|
|
- A second, non-superuser requester who may change the document
|
|
but has no permission to view that tag
|
|
WHEN:
|
|
- ai_suggestions is requested by the second requester and the
|
|
cache is hit
|
|
THEN:
|
|
- The cache hit still runs permission filtering fresh for this
|
|
requester; the invisible tag id does not leak into either the
|
|
matched or suggested tags
|
|
"""
|
|
tag_owner = User.objects.create_user(username="cache_tag_owner")
|
|
invisible_tag = Tag.objects.create(name="cache_restricted", owner=tag_owner)
|
|
requester = User.objects.create_user(username="cache_requester")
|
|
requester.user_permissions.add(
|
|
*Permission.objects.filter(
|
|
codename__in=["view_document", "change_document", "view_tag"],
|
|
),
|
|
)
|
|
mock_get_cache.return_value = MagicMock(
|
|
suggestions={
|
|
"title": "Untitled",
|
|
"tags": {"existing_ids": [invisible_tag.pk], "new_names": []},
|
|
"correspondents": {"existing_ids": [], "new_names": []},
|
|
"document_types": {"existing_ids": [], "new_names": []},
|
|
"storage_paths": {"existing_ids": [], "new_names": []},
|
|
"dates": [],
|
|
},
|
|
)
|
|
|
|
self.client.force_login(user=requester)
|
|
response = self.client.get(
|
|
f"/api/documents/{self.document.pk}/ai_suggestions/",
|
|
)
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(response.json()["tags"], [])
|
|
self.assertEqual(response.json()["suggested_tags"], [])
|
|
|
|
@patch("documents.views.get_ai_document_classification")
|
|
@override_settings(
|
|
AI_ENABLED=True,
|
|
LLM_BACKEND="mock_backend",
|
|
)
|
|
def test_ai_suggestions_with_ai_enabled(
|
|
self,
|
|
mock_get_ai_classification,
|
|
) -> None:
|
|
mock_get_ai_classification.return_value = {
|
|
"title": "AI Title",
|
|
"tags": {"existing_ids": [self.tag1.pk], "new_names": ["tag2"]},
|
|
"correspondents": {
|
|
"existing_ids": [self.correspondent1.pk],
|
|
"new_names": [],
|
|
},
|
|
"document_types": {
|
|
"existing_ids": [self.document_type1.pk],
|
|
"new_names": [],
|
|
},
|
|
"storage_paths": {"existing_ids": [self.path1.pk], "new_names": []},
|
|
"dates": ["2023-01-01"],
|
|
}
|
|
|
|
self.client.force_login(user=self.user)
|
|
response = self.client.get(
|
|
f"/api/documents/{self.document.pk}/ai_suggestions/",
|
|
)
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(
|
|
response.json(),
|
|
{
|
|
"title": "AI Title",
|
|
"tags": [self.tag1.pk],
|
|
"suggested_tags": ["tag2"],
|
|
"correspondents": [self.correspondent1.pk],
|
|
"suggested_correspondents": [],
|
|
"document_types": [self.document_type1.pk],
|
|
"suggested_document_types": [],
|
|
"storage_paths": [self.path1.pk],
|
|
"suggested_storage_paths": [],
|
|
"dates": ["2023-01-01"],
|
|
},
|
|
)
|
|
mock_get_ai_classification.assert_called_once_with(
|
|
self.document,
|
|
self.user,
|
|
None,
|
|
)
|
|
|
|
@patch("documents.views.get_ai_document_classification")
|
|
@override_settings(
|
|
AI_ENABLED=True,
|
|
LLM_BACKEND="mock_backend",
|
|
)
|
|
def test_ai_suggestions_uses_user_display_language(
|
|
self,
|
|
mock_get_ai_classification,
|
|
) -> None:
|
|
UiSettings.objects.create(user=self.user, settings={"language": "de-de"})
|
|
mock_get_ai_classification.return_value = {
|
|
"title": "KI Title",
|
|
"tags": {"existing_ids": [], "new_names": []},
|
|
"correspondents": {"existing_ids": [], "new_names": []},
|
|
"document_types": {"existing_ids": [], "new_names": []},
|
|
"storage_paths": {"existing_ids": [], "new_names": []},
|
|
"dates": [],
|
|
}
|
|
|
|
self.client.force_login(user=self.user)
|
|
response = self.client.get(
|
|
f"/api/documents/{self.document.pk}/ai_suggestions/",
|
|
)
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
mock_get_ai_classification.assert_called_once_with(
|
|
self.document,
|
|
self.user,
|
|
"de-de",
|
|
)
|
|
self.assertEqual(
|
|
get_llm_suggestion_cache(
|
|
self.document.pk,
|
|
backend="mock_backend:de-de",
|
|
).suggestions["title"],
|
|
"KI Title",
|
|
)
|
|
|
|
@patch("documents.views.get_ai_document_classification")
|
|
@override_settings(
|
|
AI_ENABLED=True,
|
|
LLM_BACKEND="mock_backend",
|
|
LLM_OUTPUT_LANGUAGE="fr-fr",
|
|
)
|
|
def test_ai_suggestions_configured_language_takes_precedence(
|
|
self,
|
|
mock_get_ai_classification,
|
|
) -> None:
|
|
UiSettings.objects.create(user=self.user, settings={"language": "de-de"})
|
|
mock_get_ai_classification.return_value = {
|
|
"title": "Titre IA",
|
|
"tags": {"existing_ids": [], "new_names": []},
|
|
"correspondents": {"existing_ids": [], "new_names": []},
|
|
"document_types": {"existing_ids": [], "new_names": []},
|
|
"storage_paths": {"existing_ids": [], "new_names": []},
|
|
"dates": [],
|
|
}
|
|
|
|
self.client.force_login(user=self.user)
|
|
response = self.client.get(
|
|
f"/api/documents/{self.document.pk}/ai_suggestions/",
|
|
)
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
mock_get_ai_classification.assert_called_once_with(
|
|
self.document,
|
|
self.user,
|
|
"fr-fr",
|
|
)
|
|
self.assertEqual(
|
|
get_llm_suggestion_cache(
|
|
self.document.pk,
|
|
backend="mock_backend:fr-fr",
|
|
).suggestions["title"],
|
|
"Titre IA",
|
|
)
|
|
|
|
@patch("documents.views.get_ai_document_classification")
|
|
@override_settings(
|
|
AI_ENABLED=True,
|
|
LLM_BACKEND="mock_backend",
|
|
LLM_MODEL="model-a",
|
|
LLM_ENDPOINT="http://endpoint-a",
|
|
)
|
|
def test_ai_suggestions_cache_key_includes_model_and_endpoint(
|
|
self,
|
|
mock_get_ai_classification,
|
|
) -> None:
|
|
"""Cached suggestions are keyed by model and endpoint, so switching
|
|
either yields a cache miss instead of a stale hit."""
|
|
mock_get_ai_classification.return_value = {
|
|
"title": "Answer A",
|
|
"tags": {"existing_ids": [], "new_names": []},
|
|
"correspondents": {"existing_ids": [], "new_names": []},
|
|
"document_types": {"existing_ids": [], "new_names": []},
|
|
"storage_paths": {"existing_ids": [], "new_names": []},
|
|
"dates": [],
|
|
}
|
|
|
|
self.client.force_login(user=self.user)
|
|
response = self.client.get(
|
|
f"/api/documents/{self.document.pk}/ai_suggestions/",
|
|
)
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
|
|
# Cached under a key that carries model + endpoint...
|
|
self.assertIsNotNone(
|
|
get_llm_suggestion_cache(
|
|
self.document.pk,
|
|
backend="mock_backend:model-a:http://endpoint-a",
|
|
),
|
|
)
|
|
|
|
@patch("documents.views.get_ai_document_classification")
|
|
@override_settings(
|
|
AI_ENABLED=True,
|
|
LLM_BACKEND="openai-like",
|
|
)
|
|
def test_ai_suggestions_with_invalid_ai_configuration(
|
|
self,
|
|
mock_get_ai_classification,
|
|
) -> None:
|
|
mock_get_ai_classification.side_effect = ValueError(
|
|
"Unknown model 'gpt-5.4-mini-2026-03-17'.",
|
|
)
|
|
|
|
self.client.force_login(user=self.user)
|
|
response = self.client.get(
|
|
f"/api/documents/{self.document.pk}/ai_suggestions/",
|
|
)
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
|
|
self.assertEqual(
|
|
response.json(),
|
|
{
|
|
"ai": ["Invalid AI configuration."],
|
|
},
|
|
)
|
|
self.assertIsNone(
|
|
get_llm_suggestion_cache(self.document.pk, backend="openai-like"),
|
|
)
|
|
|
|
@patch("documents.views.get_ai_document_classification")
|
|
@override_settings(
|
|
AI_ENABLED=True,
|
|
LLM_BACKEND="openai-like",
|
|
)
|
|
def test_ai_suggestions_with_llm_timeout(
|
|
self,
|
|
mock_get_ai_classification,
|
|
) -> None:
|
|
mock_get_ai_classification.side_effect = LLMTimeoutError()
|
|
|
|
self.client.force_login(user=self.user)
|
|
response = self.client.get(
|
|
f"/api/documents/{self.document.pk}/ai_suggestions/",
|
|
)
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_503_SERVICE_UNAVAILABLE)
|
|
self.assertEqual(
|
|
response.json(),
|
|
{
|
|
"ai": ["AI backend request timed out."],
|
|
},
|
|
)
|
|
self.assertIsNone(
|
|
get_llm_suggestion_cache(self.document.pk, backend="openai-like"),
|
|
)
|
|
|
|
@patch("documents.views.get_ai_document_classification")
|
|
@override_settings(
|
|
AI_ENABLED=True,
|
|
LLM_BACKEND="mock_backend",
|
|
)
|
|
def test_ai_suggestions_combines_existing_ids_and_new_names(
|
|
self,
|
|
mock_get_ai_classification,
|
|
) -> None:
|
|
"""
|
|
GIVEN:
|
|
- AI classification returns a taxonomy choice with both an
|
|
existing tag id and a new tag name not present in the database
|
|
WHEN:
|
|
- ai_suggestions is requested
|
|
THEN:
|
|
- the existing id is resolved into the matched tags list
|
|
- the new name is fuzzy-matched, and since it doesn't match any
|
|
existing tag, it is surfaced as a suggested tag
|
|
"""
|
|
mock_get_ai_classification.return_value = {
|
|
"title": "Lab Report",
|
|
"tags": {"existing_ids": [self.tag1.pk], "new_names": ["Follow-up"]},
|
|
"correspondents": {"existing_ids": [], "new_names": []},
|
|
"document_types": {"existing_ids": [], "new_names": []},
|
|
"storage_paths": {"existing_ids": [], "new_names": []},
|
|
"dates": [],
|
|
}
|
|
|
|
self.client.force_login(user=self.user)
|
|
response = self.client.get(
|
|
f"/api/documents/{self.document.pk}/ai_suggestions/",
|
|
)
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(response.json()["tags"], [self.tag1.pk])
|
|
self.assertEqual(response.json()["suggested_tags"], ["Follow-up"])
|
|
|
|
@patch("documents.views.get_ai_document_classification")
|
|
@override_settings(
|
|
AI_ENABLED=True,
|
|
LLM_BACKEND="mock_backend",
|
|
)
|
|
def test_ai_suggestions_deduplicates_id_matched_via_both_paths(
|
|
self,
|
|
mock_get_ai_classification,
|
|
) -> None:
|
|
"""
|
|
GIVEN:
|
|
- AI classification returns the same tag both as an existing_id
|
|
and as a new_name that fuzzy-matches that same tag
|
|
WHEN:
|
|
- ai_suggestions is requested
|
|
THEN:
|
|
- The tag's id appears exactly once in the response, not twice
|
|
"""
|
|
mock_get_ai_classification.return_value = {
|
|
"title": "Lab Report",
|
|
"tags": {
|
|
"existing_ids": [self.tag1.pk],
|
|
"new_names": [self.tag1.name],
|
|
},
|
|
"correspondents": {"existing_ids": [], "new_names": []},
|
|
"document_types": {"existing_ids": [], "new_names": []},
|
|
"storage_paths": {"existing_ids": [], "new_names": []},
|
|
"dates": [],
|
|
}
|
|
|
|
self.client.force_login(user=self.user)
|
|
response = self.client.get(
|
|
f"/api/documents/{self.document.pk}/ai_suggestions/",
|
|
)
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(response.json()["tags"], [self.tag1.pk])
|
|
self.assertEqual(response.json()["suggested_tags"], [])
|
|
|
|
@patch("documents.views.get_ai_document_classification")
|
|
@override_settings(
|
|
AI_ENABLED=True,
|
|
LLM_BACKEND="mock_backend",
|
|
)
|
|
def test_ai_suggestions_existing_id_not_visible_falls_through_to_suggested(
|
|
self,
|
|
mock_get_ai_classification,
|
|
) -> None:
|
|
"""
|
|
GIVEN:
|
|
- A non-superuser who may change the document but has no
|
|
permission to view a tag owned by somebody else
|
|
- AI classification returns that tag's id in existing_ids (e.g.
|
|
from a cached response generated for a broader-visibility user)
|
|
WHEN:
|
|
- ai_suggestions is requested by that user
|
|
THEN:
|
|
- the invisible id is silently dropped by resolve_tag_ids, so
|
|
permission filtering survives the full request path
|
|
- it does not appear in either the matched or suggested tags
|
|
"""
|
|
tag_owner = User.objects.create_user(username="tagowner")
|
|
invisible_tag = Tag.objects.create(name="restricted", owner=tag_owner)
|
|
requester = User.objects.create_user(username="requester")
|
|
requester.user_permissions.add(
|
|
*Permission.objects.filter(
|
|
codename__in=["view_document", "change_document", "view_tag"],
|
|
),
|
|
)
|
|
|
|
mock_get_ai_classification.return_value = {
|
|
"title": "Untitled",
|
|
"tags": {"existing_ids": [invisible_tag.pk], "new_names": []},
|
|
"correspondents": {"existing_ids": [], "new_names": []},
|
|
"document_types": {"existing_ids": [], "new_names": []},
|
|
"storage_paths": {"existing_ids": [], "new_names": []},
|
|
"dates": [],
|
|
}
|
|
|
|
self.client.force_login(user=requester)
|
|
response = self.client.get(
|
|
f"/api/documents/{self.document.pk}/ai_suggestions/",
|
|
)
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
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)
|
|
suggestions = {
|
|
"title": "AI Title",
|
|
"tags": ["tag1", "tag2"],
|
|
"correspondents": ["correspondent1"],
|
|
"document_types": ["type1"],
|
|
"storage_paths": ["path1"],
|
|
"dates": ["2023-01-01"],
|
|
}
|
|
set_llm_suggestions_cache(
|
|
self.document.pk,
|
|
suggestions,
|
|
backend="mock_backend",
|
|
)
|
|
self.assertEqual(
|
|
get_llm_suggestion_cache(
|
|
self.document.pk,
|
|
backend="mock_backend",
|
|
).suggestions,
|
|
suggestions,
|
|
)
|
|
# post_save signal triggered
|
|
update_llm_suggestions_cache(
|
|
sender=None,
|
|
instance=self.document,
|
|
)
|
|
self.assertIsNone(
|
|
get_llm_suggestion_cache(
|
|
self.document.pk,
|
|
backend="mock_backend",
|
|
),
|
|
)
|
|
|
|
|
|
class TestAIChatStreamingView(DirectoriesMixin, TestCase):
|
|
ENDPOINT = "/api/documents/chat/"
|
|
|
|
def setUp(self) -> None:
|
|
self.user = User.objects.create_user(username="testuser", password="pass")
|
|
self.client.force_login(user=self.user)
|
|
self.document = Document.objects.create(
|
|
title="Test Document",
|
|
filename="test.pdf",
|
|
mime_type="application/pdf",
|
|
)
|
|
super().setUp()
|
|
|
|
def grant_view_document_permission(self) -> None:
|
|
self.user.user_permissions.add(
|
|
*Permission.objects.filter(codename="view_document"),
|
|
)
|
|
|
|
@override_settings(AI_ENABLED=False)
|
|
def test_post_ai_disabled(self) -> None:
|
|
self.grant_view_document_permission()
|
|
response = self.client.post(
|
|
self.ENDPOINT,
|
|
data='{"q": "question"}',
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(response.status_code, 400)
|
|
self.assertIn(b"AI is required for this feature", response.content)
|
|
|
|
@patch("documents.views.stream_chat_with_documents")
|
|
@patch("documents.views.permitted_document_ids")
|
|
@override_settings(AI_ENABLED=True)
|
|
def test_post_no_document_id(self, mock_permitted_ids, mock_stream_chat) -> None:
|
|
self.grant_view_document_permission()
|
|
mock_permitted_ids.return_value = [self.document.pk]
|
|
mock_stream_chat.return_value = iter([b"data"])
|
|
response = self.client.post(
|
|
self.ENDPOINT,
|
|
data='{"q": "question"}',
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertEqual(response["Content-Type"], "text/event-stream")
|
|
mock_stream_chat.assert_called_once()
|
|
call_kwargs = mock_stream_chat.call_args.kwargs
|
|
self.assertEqual(call_kwargs["query_str"], "question")
|
|
self.assertEqual(list(call_kwargs["documents"]), [self.document])
|
|
self.assertIsNone(call_kwargs["output_language"])
|
|
|
|
@patch("documents.views.stream_chat_with_documents")
|
|
@patch("documents.views.permitted_document_ids")
|
|
@override_settings(AI_ENABLED=True)
|
|
def test_post_uses_user_display_language(
|
|
self,
|
|
mock_permitted_ids,
|
|
mock_stream_chat,
|
|
) -> None:
|
|
UiSettings.objects.create(user=self.user, settings={"language": "de-de"})
|
|
self.grant_view_document_permission()
|
|
mock_permitted_ids.return_value = [self.document.pk]
|
|
mock_stream_chat.return_value = iter([b"data"])
|
|
|
|
response = self.client.post(
|
|
self.ENDPOINT,
|
|
data='{"q": "question"}',
|
|
content_type="application/json",
|
|
)
|
|
|
|
self.assertEqual(response.status_code, 200)
|
|
mock_stream_chat.assert_called_once()
|
|
call_kwargs = mock_stream_chat.call_args.kwargs
|
|
self.assertEqual(call_kwargs["query_str"], "question")
|
|
self.assertEqual(list(call_kwargs["documents"]), [self.document])
|
|
self.assertEqual(call_kwargs["output_language"], "de-de")
|
|
|
|
@patch("documents.views.stream_chat_with_documents")
|
|
@override_settings(AI_ENABLED=True)
|
|
def test_post_with_document_id(self, mock_stream_chat) -> None:
|
|
self.grant_view_document_permission()
|
|
mock_stream_chat.return_value = iter([b"data"])
|
|
response = self.client.post(
|
|
self.ENDPOINT,
|
|
data=f'{{"q": "question", "document_id": {self.document.pk}}}',
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertEqual(response["Content-Type"], "text/event-stream")
|
|
|
|
@override_settings(AI_ENABLED=True)
|
|
def test_post_with_invalid_document_id(self) -> None:
|
|
self.grant_view_document_permission()
|
|
response = self.client.post(
|
|
self.ENDPOINT,
|
|
data='{"q": "question", "document_id": 999999}',
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(response.status_code, 400)
|
|
self.assertIn(b"Document not found", response.content)
|
|
|
|
@patch("documents.views.has_perms_owner_aware")
|
|
@override_settings(AI_ENABLED=True)
|
|
def test_post_with_document_id_no_permission(self, mock_has_perms) -> None:
|
|
self.grant_view_document_permission()
|
|
mock_has_perms.return_value = False
|
|
response = self.client.post(
|
|
self.ENDPOINT,
|
|
data=f'{{"q": "question", "document_id": {self.document.pk}}}',
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(response.status_code, 403)
|
|
self.assertIn(b"Insufficient permissions", response.content)
|
|
|
|
@patch("documents.views.stream_chat_with_documents")
|
|
@override_settings(AI_ENABLED=True)
|
|
def test_post_no_document_id_requires_view_document_permission(
|
|
self,
|
|
mock_stream_chat,
|
|
) -> None:
|
|
response = self.client.post(
|
|
self.ENDPOINT,
|
|
data='{"q": "question"}',
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(response.status_code, 403)
|
|
mock_stream_chat.assert_not_called()
|
|
|
|
@patch("documents.views.stream_chat_with_documents")
|
|
@override_settings(AI_ENABLED=True)
|
|
def test_post_with_document_id_requires_view_document_permission(
|
|
self,
|
|
mock_stream_chat,
|
|
) -> None:
|
|
response = self.client.post(
|
|
self.ENDPOINT,
|
|
data=f'{{"q": "question", "document_id": {self.document.pk}}}',
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(response.status_code, 403)
|
|
mock_stream_chat.assert_not_called()
|