Compare commits

..
Author SHA1 Message Date
stumpylog 16f1427b7a Not bad catches from Copilot, if a little extra secure 2026-08-27 11:37:07 -07:00
stumpylog 8f008a8bf4 feature: add Tantivy full-text fallback adapter for taxonomy candidates
This brings users without an embedding backend configured to closer
parity with those who do.  Reuse the search backend to locate similar
documents and use them to provide the LLM with the better suggestion pool
to draw from
2026-08-27 09:42:08 -07:00
6 changed files with 449 additions and 383 deletions
+33 -64
View File
@@ -305,74 +305,46 @@ def modify_custom_fields(
else [(field, None) for field in add_custom_fields]
)
custom_fields_by_id: dict[int, CustomField] = {
cf.id: cf
for cf in CustomField.objects.filter(
id__in=[int(field) for field, _ in add_custom_fields],
)
}
# Deferred, not `.only()`: signal receivers touch other Document fields,
# and `.only("pk")` would just turn that into a per-document reload.
# `content` is the one field both large and unused here. Skipped
# entirely for a remove-only call -- the removal pass below resolves
# its own documents.
docs_by_id: dict[int, Document] = (
{
doc.id: doc
for doc in Document.objects.filter(id__in=affected_docs).defer("content")
}
if add_custom_fields
else {}
)
custom_fields = CustomField.objects.filter(
id__in=[int(field) for field, _ in add_custom_fields],
).distinct()
for field_id, value in add_custom_fields:
custom_field = custom_fields_by_id[field_id]
value_field = CustomFieldInstance.TYPE_TO_DATA_STORE_NAME_MAP[
custom_field.data_type
]
for doc_id in affected_docs:
defaults = {value_field: value}
if (
custom_field.data_type == CustomField.FieldDataType.DOCUMENTLINK
and value
and doc_id in value
):
# Prevent self-linking
continue
# Not update_or_create(): it fetches an existing row via plain
# `.get()` before calling .save(), so a signal receiver touching
# `.field`/`.document` (e.g. auditlog) on that save re-fetches
# per instance regardless of what's passed in as lookup kwargs.
# Assigning the cached objects ourselves before .save() avoids
# that for both the create and update case.
try:
instance = CustomFieldInstance.objects.get(
document=docs_by_id[doc_id],
field=custom_field,
)
except CustomFieldInstance.DoesNotExist:
instance = CustomFieldInstance(
document=docs_by_id[doc_id],
field=custom_field,
)
instance.document = docs_by_id[doc_id]
instance.field = custom_field
for attr, val in defaults.items():
setattr(instance, attr, val)
instance.save()
defaults = {}
custom_field = custom_fields.get(id=field_id)
if custom_field:
value_field = CustomFieldInstance.TYPE_TO_DATA_STORE_NAME_MAP[
custom_field.data_type
]
defaults[value_field] = value
if (
custom_field.data_type == CustomField.FieldDataType.DOCUMENTLINK
and value
and doc_id in value
):
# Prevent self-linking
continue
CustomFieldInstance.objects.update_or_create(
document_id=doc_id,
field_id=field_id,
defaults=defaults,
)
if custom_field.data_type == CustomField.FieldDataType.DOCUMENTLINK:
reflect_doclinks(docs_by_id[doc_id], custom_field, value)
doc = Document.objects.get(id=doc_id)
reflect_doclinks(doc, custom_field, value)
# For doc link fields being removed, remove symmetrical links.
# select_related here avoids resolving every affected document up front.
# For doc link fields that are being removed, remove symmetrical links
for doclink_being_removed_instance in CustomFieldInstance.objects.filter(
document_id__in=affected_docs,
field__id__in=remove_custom_fields,
field__data_type=CustomField.FieldDataType.DOCUMENTLINK,
value_document_ids__isnull=False,
).select_related("field", "document"):
):
for target_doc_id in doclink_being_removed_instance.value:
remove_doclink(
document=doclink_being_removed_instance.document,
document=Document.objects.get(
id=doclink_being_removed_instance.document.id,
),
field=doclink_being_removed_instance.field,
target_doc_id=target_doc_id,
)
@@ -1205,13 +1177,10 @@ def remove_doclink(
"""
Removes a 'symmetrical' link to `document` from the target document's existing custom field instance
"""
# select_related: a signal receiver (auditlog) touches .document/.field
# on save() below -- without this, that's a per-call reload query.
target_doc_field_instance = (
CustomFieldInstance.objects.filter(document_id=target_doc_id, field=field)
.select_related("document", "field")
.first()
)
target_doc_field_instance = CustomFieldInstance.objects.filter(
document_id=target_doc_id,
field=field,
).first()
if (
target_doc_field_instance is not None
and document.id in target_doc_field_instance.value
-198
View File
@@ -6,9 +6,7 @@ from unittest import mock
import pikepdf
from django.contrib.auth.models import Group
from django.contrib.auth.models import User
from django.db import connection
from django.test import TestCase
from django.test.utils import CaptureQueriesContext
from guardian.shortcuts import assign_perm
from guardian.shortcuts import get_groups_with_perms
from guardian.shortcuts import get_users_with_perms
@@ -346,202 +344,6 @@ class TestBulkEdit(DirectoriesMixin, TestCase):
assert _cf_3 is not None
self.assertNotIn(self.doc3.id, _cf_3.value)
def test_modify_custom_fields_batches_field_lookup(self) -> None:
"""
GIVEN:
- Several documents are being bulk-edited to add several custom
fields at once
WHEN:
- modify_custom_fields runs
THEN:
- Each CustomField is resolved with one batched query total, not
once per (field, document) pair
"""
docs = [
Document.objects.create(checksum=f"batch-{i}", title=f"batch-{i}")
for i in range(6)
]
fields = [
CustomField.objects.create(
name=f"Batch Field {i}",
data_type=CustomField.FieldDataType.STRING,
)
for i in range(4)
]
with CaptureQueriesContext(connection) as ctx:
bulk_edit.modify_custom_fields(
[doc.id for doc in docs],
add_custom_fields=[field.id for field in fields],
remove_custom_fields=[],
)
field_lookups = [
q
for q in ctx.captured_queries
if 'FROM "documents_customfield"' in q["sql"]
]
self.assertEqual(
len(field_lookups),
1,
"Expected a single batched query to resolve the custom fields, "
f"got {len(field_lookups)}: {field_lookups}",
)
for doc in docs:
self.assertEqual(doc.custom_fields.count(), len(fields))
def test_modify_custom_fields_batches_document_lookup_for_documentlink(
self,
) -> None:
"""
GIVEN:
- Several documents are being bulk-edited to add a DOCUMENTLINK
custom field at once
WHEN:
- modify_custom_fields runs
THEN:
- The Document rows needed to reflect the symmetrical links are
resolved with one batched query total, not once per document
"""
docs = [
Document.objects.create(checksum=f"link-{i}", title=f"link-{i}")
for i in range(6)
]
target = Document.objects.create(checksum="link-target", title="link-target")
doclink_field = CustomField.objects.create(
name="Related",
data_type=CustomField.FieldDataType.DOCUMENTLINK,
)
with CaptureQueriesContext(connection) as ctx:
bulk_edit.modify_custom_fields(
[doc.id for doc in docs],
add_custom_fields={doclink_field.id: [target.id]},
remove_custom_fields=[],
)
single_document_lookups = [
q
for q in ctx.captured_queries
if 'FROM "documents_document"' in q["sql"]
and '"documents_document"."id" = ' in q["sql"]
]
self.assertEqual(
len(single_document_lookups),
0,
"Expected document rows to come from a batched query, not "
f"per-document lookups, got: {single_document_lookups}",
)
for doc in docs:
self.assertEqual(
doc.custom_fields.get(field=doclink_field).value,
[target.id],
)
def test_modify_custom_fields_update_caches_document_and_field(self) -> None:
"""
GIVEN:
- Several documents already have an instance of a custom field
WHEN:
- modify_custom_fields runs again for the same field, updating
the existing instances rather than creating new ones
THEN:
- No per-instance `.document`/`.field` reload query is issued
(e.g. by auditlog's post_save receiver touching them)
"""
docs = [
Document.objects.create(checksum=f"update-{i}", title=f"update-{i}")
for i in range(6)
]
field = CustomField.objects.create(
name="Update Field",
data_type=CustomField.FieldDataType.STRING,
)
bulk_edit.modify_custom_fields(
[doc.id for doc in docs],
add_custom_fields=[field.id],
remove_custom_fields=[],
)
with CaptureQueriesContext(connection) as ctx:
bulk_edit.modify_custom_fields(
[doc.id for doc in docs],
add_custom_fields={field.id: "updated value"},
remove_custom_fields=[],
)
single_row_reloads = [
q
for q in ctx.captured_queries
if ('FROM "documents_document"' in q["sql"] and '."id" = ' in q["sql"])
or ('FROM "documents_customfield"' in q["sql"] and '."id" = ' in q["sql"])
]
self.assertEqual(
single_row_reloads,
[],
"Expected no per-instance document/field reload queries when "
f"updating existing custom field instances, got: {single_row_reloads}",
)
for doc in docs:
self.assertEqual(
doc.custom_fields.get(field=field).value,
"updated value",
)
def test_modify_custom_fields_removes_symmetrical_doclinks_batched(self) -> None:
"""
GIVEN:
- Several source documents link to a shared target via a doc
link field
WHEN:
- The field is removed from all of them in one call
THEN:
- The symmetrical links are removed from the target
- No per-document lookup query is issued, on either side
"""
target = Document.objects.create(checksum="rm-target", title="rm-target")
docs = [
Document.objects.create(checksum=f"rm-{i}", title=f"rm-{i}")
for i in range(6)
]
field = CustomField.objects.create(
name="Related",
data_type=CustomField.FieldDataType.DOCUMENTLINK,
)
bulk_edit.modify_custom_fields(
[doc.id for doc in docs],
add_custom_fields={field.id: [target.id]},
remove_custom_fields=[],
)
self.assertEqual(
target.custom_fields.get(field=field).value,
[d.id for d in docs],
)
with CaptureQueriesContext(connection) as ctx:
bulk_edit.modify_custom_fields(
[doc.id for doc in docs],
add_custom_fields=[],
remove_custom_fields=[field.id],
)
single_document_lookups = [
q
for q in ctx.captured_queries
if 'FROM "documents_document"' in q["sql"]
and '"documents_document"."id" = ' in q["sql"]
]
self.assertEqual(
single_document_lookups,
[],
"Expected batched document resolution, not per-document lookups, "
f"got: {single_document_lookups}",
)
self.assertEqual(target.custom_fields.get(field=field).value, [])
def test_modify_custom_fields_doclink_self_link(self) -> None:
"""
GIVEN:
+99 -54
View File
@@ -5,13 +5,14 @@ from django.conf import settings
from django.contrib.auth.models import User
from documents.models import Document
from documents.permissions import get_objects_for_user_owner_aware
from documents.permissions import permitted_object_ids
from documents.permissions import restrict_queryset_to_visible
from documents.permissions import user_is_unrestricted
from paperless.config import AIConfig
from paperless_ai.base_model import ClassificationSuggestions
from paperless_ai.base_model import TaxonomyChoiceDict
from paperless_ai.client import AIClient
from paperless_ai.db import db_connection_released
from paperless_ai.indexing import _node_document_ids
from paperless_ai.indexing import retrieve_similar_nodes
from paperless_ai.indexing import truncate_content
from paperless_ai.prompts.context import ClassificationPromptContext
@@ -19,7 +20,9 @@ from paperless_ai.prompts.context import LocalizationPromptContext
from paperless_ai.prompts.context import RagContextPromptContext
from paperless_ai.prompts.render import render_prompt
from paperless_ai.taxonomy import AssignedMetadata
from paperless_ai.taxonomy import SimilarDocument
from paperless_ai.taxonomy import TaxonomyCandidates
from paperless_ai.taxonomy import _node_document_weights
from paperless_ai.taxonomy import build_taxonomy_candidates
from paperless_ai.taxonomy import empty_taxonomy_candidates
from paperless_ai.taxonomy import format_taxonomy_for_prompt
@@ -39,6 +42,48 @@ logger = logging.getLogger("paperless_ai.rag_classifier")
TAXONOMY_CANDIDATE_TOP_K = 15
def _fulltext_similar_documents(
document: Document,
user: User | None,
top_k: int,
) -> list[SimilarDocument]:
"""Rank-based fallback when no embedding backend is configured. Uses
Tantivy's "More Like This" (term-overlap similarity) instead of vector
similarity - cruder, but far better than no candidates at all.
more_like_this_ids returns only a ranked ID list, no scores, so weight is
synthesized from rank (descending from top_k) rather than claiming a
similarity magnitude that doesn't exist. An unrestricted user (none, or an
active superuser - see user_is_unrestricted) is normalized to ``None``
before calling, since the backend's permission filter has no superuser
short-circuit of its own. Results are re-checked with
restrict_queryset_to_visible() since Tantivy's indexed permission fields
lag the DB via async reindexing.
"""
from documents.search import get_backend
unrestricted = user_is_unrestricted(user)
search_user = None if unrestricted else user
backend = get_backend()
similar_ids = backend.more_like_this_ids(
document.pk,
user=search_user,
limit=top_k,
)
if not unrestricted:
allowed_ids = set(
restrict_queryset_to_visible(
Document.objects.filter(pk__in=similar_ids),
user,
"view_document",
).values_list("pk", flat=True),
)
similar_ids = [doc_id for doc_id in similar_ids if doc_id in allowed_ids]
return [
SimilarDocument(document_id=doc_id, weight=float(top_k - rank))
for rank, doc_id in enumerate(similar_ids)
]
def get_language_name(language_code: str) -> str:
normalized_language_code = language_code.lower()
for code, name in settings.LANGUAGES:
@@ -147,45 +192,54 @@ def get_taxonomy_context(
user: User | None = None,
max_docs: int = 5,
) -> tuple[TaxonomyCandidates, AssignedMetadata, str]:
"""One retrieval feeds both taxonomy candidates and RAG text context.
On any retrieval failure, degrades to empty candidates/context rather than
propagating the exception - a vector-store outage should not block
classification, only its RAG-assisted enrichment.
"""One retrieval feeds both taxonomy candidates and RAG text context. Uses
vector similarity when an embedding backend is configured, otherwise
falls back to Tantivy full-text "More Like This" similarity - see
_fulltext_similar_documents. On any retrieval failure, degrades to empty
candidates/context rather than propagating the exception - neither a
vector-store outage nor a search-index issue should block classification,
only its context-assisted enrichment.
"""
assigned = get_assigned_metadata(document, user)
ai_config = AIConfig()
try:
# None means "no restriction" to retrieve_similar_nodes. A superuser
# (like no user at all) can see every document, so skip materializing
# every visible pk into a Python list and passing it through as an IN
# filter: for a large library that is a wasted quadratic scan in the
# vector store at best, and past ~32,763 documents a hard
# sqlite3.OperationalError (SQLite's bound-parameter limit) at worst.
# get_objects_for_user_owner_aware() would return every Document for a
# superuser anyway (guardian's own with_superuser shortcut), so this
# changes nothing about which documents are considered -- only how we
# get there.
visible_document_ids = (
None
if user is None or user.is_superuser
else list(
get_objects_for_user_owner_aware(
user,
"view_document",
Document,
).values_list("pk", flat=True),
if ai_config.llm_embedding_backend:
# None means "no restriction" to retrieve_similar_nodes. A superuser
# (like no user at all) can see every document, so skip materializing
# every visible pk into a Python list and passing it through as an IN
# filter: for a large library that is a wasted quadratic scan in the
# vector store at best, and past ~32,763 documents a hard
# sqlite3.OperationalError (SQLite's bound-parameter limit) at worst.
# permitted_object_ids() has its own superuser shortcut that would
# return every Document's id anyway, so this changes nothing about
# which documents are considered -- only how we get there.
visible_document_ids = (
None
if user is None or user.is_superuser
else list(permitted_object_ids(user, Document, "view_document"))
)
nodes = retrieve_similar_nodes(
document,
top_k=TAXONOMY_CANDIDATE_TOP_K,
document_ids=visible_document_ids,
)
similar_documents = _node_document_weights(nodes)
else:
# See _fulltext_similar_documents: it applies its own permission
# filter via `user`, so no visible-document-id list is needed here.
similar_documents = _fulltext_similar_documents(
document,
user,
top_k=TAXONOMY_CANDIDATE_TOP_K,
)
)
nodes = retrieve_similar_nodes(
document,
top_k=TAXONOMY_CANDIDATE_TOP_K,
document_ids=visible_document_ids,
)
candidates = build_taxonomy_candidates(nodes, user)
candidates = build_taxonomy_candidates(similar_documents, user)
similar_docs = list(
Document.objects.filter(pk__in=_node_document_ids(nodes))[:max_docs],
)
similar_doc_ids = [s["document_id"] for s in similar_documents]
docs_by_id = Document.objects.in_bulk(similar_doc_ids)
similar_docs = [
docs_by_id[doc_id] for doc_id in similar_doc_ids if doc_id in docs_by_id
][:max_docs]
context_blocks = []
for similar in similar_docs:
text = similar.content[:1000] or ""
@@ -193,8 +247,8 @@ def get_taxonomy_context(
context_blocks.append(f"TITLE: {title}\n{text}")
except Exception:
logger.exception(
"Failed to retrieve RAG neighbours for document %s; continuing "
"without taxonomy candidates or similar-document context.",
"Failed to retrieve similar-document context for document %s; "
"continuing without taxonomy candidates or similar-document context.",
document.pk,
)
return empty_taxonomy_candidates(), assigned, ""
@@ -277,23 +331,14 @@ def get_ai_document_classification(
) -> ClassificationSuggestions:
ai_config = AIConfig()
if ai_config.llm_embedding_backend:
candidates, assigned, context = get_taxonomy_context(document, user)
prompt = build_prompt_with_rag(
document,
ai_config,
candidates=candidates,
assigned=assigned,
context=context,
)
else:
candidates = empty_taxonomy_candidates()
prompt = build_prompt_without_rag(
document,
ai_config,
candidates=candidates,
assigned=get_assigned_metadata(document, user),
)
candidates, assigned, context = get_taxonomy_context(document, user)
prompt = build_prompt_with_rag(
document,
ai_config,
candidates=candidates,
assigned=assigned,
context=context,
)
client = AIClient()
# Hand the pooled DB connection back while the (slow) LLM query runs so it
+27 -14
View File
@@ -33,6 +33,11 @@ class TaxonomyCandidate(TypedDict):
weight: float
class SimilarDocument(TypedDict):
document_id: int
weight: float
class TaxonomyCandidates(TypedDict):
tags: list[TaxonomyCandidate]
document_types: list[TaxonomyCandidate]
@@ -105,10 +110,10 @@ def get_assigned_metadata(document: Document, user: User | None) -> AssignedMeta
)
def _node_document_weights(nodes: list["NodeWithScore"]) -> dict[int, float]:
"""document_id -> that node's similarity score, summed if a document_id
appears more than once across the retrieved nodes (e.g. multiple chunks
of the same source document)."""
def _node_document_weights(nodes: list["NodeWithScore"]) -> list[SimilarDocument]:
"""Sum each node's similarity score into its document_id (a document can
appear via multiple chunks/nodes) and return one SimilarDocument per
distinct document_id."""
weights: dict[int, float] = defaultdict(float)
for node in nodes:
document_id = node.metadata.get("document_id")
@@ -121,7 +126,10 @@ def _node_document_weights(nodes: list["NodeWithScore"]) -> dict[int, float]:
weights[int(document_id)] += float(node.score or 0.0)
except (TypeError, ValueError): # pragma: no cover
continue
return weights
return [
SimilarDocument(document_id=document_id, weight=weight)
for document_id, weight in weights.items()
]
def _visible_ranked_candidates(
@@ -157,21 +165,26 @@ def _visible_ranked_candidates(
def build_taxonomy_candidates(
nodes: list["NodeWithScore"],
similar_documents: list[SimilarDocument],
user: User | None,
) -> TaxonomyCandidates:
"""Resolve each neighbour node's document_id to a live Document, read its
*current* tags/type/correspondent/storage_path via the ORM (never the
possibly-stale names cached in vector-index node metadata), weight each
distinct taxonomy object by aggregate neighbour similarity, permission-filter
"""Resolve each similar document's id to a live Document, read its
*current* tags/type/correspondent/storage_path via the ORM (never any
possibly-stale names an adapter's source might have cached), weight each
distinct taxonomy object by aggregate similarity weight, permission-filter
against what ``user`` can see, and return each category ranked by weight
and capped.
and capped. ``similar_documents`` may come from either the vector-RAG
adapter or the full-text fallback adapter - both produce this same shape.
"""
document_weights = _node_document_weights(nodes)
if not document_weights:
if not similar_documents:
return empty_taxonomy_candidates()
# Both adapters guarantee at most one SimilarDocument per document_id, so
# this never silently drops a duplicate's weight.
document_weights: dict[int, float] = {
s["document_id"]: s["weight"] for s in similar_documents
}
# Only .tags.all() needs prefetching (a reverse M2M, one extra query for
# the whole batch). document_type/correspondent/storage_path are read
# below via their *_id columns (neighbour.document_type_id, etc.), which
+257 -22
View File
@@ -1,3 +1,4 @@
from collections.abc import Generator
from types import SimpleNamespace
from unittest.mock import MagicMock
from unittest.mock import patch
@@ -7,10 +8,13 @@ import pytest_mock
from django.test import override_settings
from documents.models import Document
from documents.search import TantivyBackend
from documents.tests.factories import DocumentFactory
from documents.tests.factories import TagFactory
from documents.tests.factories import UserFactory
from paperless.config import AIConfig
from paperless_ai.ai_classifier import TAXONOMY_CANDIDATE_TOP_K
from paperless_ai.ai_classifier import _fulltext_similar_documents
from paperless_ai.ai_classifier import _restrict_to_shown_candidates
from paperless_ai.ai_classifier import build_localization_prompt
from paperless_ai.ai_classifier import build_prompt_with_rag
@@ -20,6 +24,7 @@ from paperless_ai.ai_classifier import get_language_name
from paperless_ai.ai_classifier import get_taxonomy_context
from paperless_ai.base_model import ClassificationSuggestions
from paperless_ai.base_model import TaxonomyChoiceDict
from paperless_ai.taxonomy import SimilarDocument
from paperless_ai.taxonomy import TaxonomyCandidate
from paperless_ai.taxonomy import TaxonomyCandidates
from paperless_ai.taxonomy import empty_taxonomy_candidates
@@ -204,12 +209,10 @@ def test_use_rag_if_configured(
@pytest.mark.django_db
@patch("paperless_ai.client.AIClient.run_llm_query")
@patch("paperless_ai.ai_classifier.build_prompt_without_rag")
@patch("paperless_ai.ai_classifier.AIConfig")
@patch("paperless_ai.ai_classifier.build_prompt_with_rag")
@override_settings(LLM_BACKEND="ollama", LLM_MODEL="some_model")
def test_use_without_rag_if_not_configured(
mock_ai_config,
mock_build_prompt_without_rag,
def test_use_rag_prompt_even_without_embedding_backend(
mock_build_prompt_with_rag,
mock_run_llm_query,
mock_document,
):
@@ -219,13 +222,13 @@ def test_use_without_rag_if_not_configured(
WHEN:
- get_ai_document_classification() is called
THEN:
- The non-RAG prompt builder is used
- The RAG-context prompt builder is still used (fed by the full-text
fallback's context/candidates instead of the vector store's)
"""
mock_ai_config.return_value.llm_embedding_backend = None
mock_build_prompt_without_rag.return_value = "Prompt without RAG"
mock_build_prompt_with_rag.return_value = "Prompt with RAG"
mock_run_llm_query.return_value = NESTED_SUGGESTIONS
get_ai_document_classification(mock_document)
mock_build_prompt_without_rag.assert_called_once()
mock_build_prompt_with_rag.assert_called_once()
@pytest.mark.django_db
@@ -303,6 +306,7 @@ def test_build_localization_prompt_preserves_unicode_characters():
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_get_taxonomy_context_assembles_rag_text_and_candidates():
"""
GIVEN:
@@ -344,6 +348,7 @@ def test_get_taxonomy_context_assembles_rag_text_and_candidates():
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_get_taxonomy_context_no_similar_docs():
"""
GIVEN:
@@ -367,6 +372,67 @@ def test_get_taxonomy_context_no_similar_docs():
}
@pytest.mark.django_db
def test_get_taxonomy_context_uses_fulltext_fallback_when_no_embedding_backend(
mocker: pytest_mock.MockerFixture,
) -> None:
"""
GIVEN:
- No LLM embedding backend is configured (the default test settings)
WHEN:
- get_taxonomy_context() is called
THEN:
- _fulltext_similar_documents() is called with the document, the user
and TAXONOMY_CANDIDATE_TOP_K
- retrieve_similar_nodes() (the vector path) is never called
"""
document = DocumentFactory.create(content="Some content")
mock_fulltext = mocker.patch(
"paperless_ai.ai_classifier._fulltext_similar_documents",
return_value=[],
)
mock_retrieve = mocker.patch("paperless_ai.ai_classifier.retrieve_similar_nodes")
get_taxonomy_context(document, user=None)
mock_fulltext.assert_called_once_with(
document,
None,
top_k=TAXONOMY_CANDIDATE_TOP_K,
)
mock_retrieve.assert_not_called()
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_get_taxonomy_context_uses_vector_path_when_embedding_backend_configured(
mocker: pytest_mock.MockerFixture,
) -> None:
"""
GIVEN:
- An LLM embedding backend is configured
WHEN:
- get_taxonomy_context() is called
THEN:
- retrieve_similar_nodes() (the vector path) is called
- _fulltext_similar_documents() (the no-embedding-backend fallback)
is never called
"""
document = DocumentFactory.create(content="Some content")
mock_retrieve = mocker.patch(
"paperless_ai.ai_classifier.retrieve_similar_nodes",
return_value=[],
)
mock_fulltext = mocker.patch(
"paperless_ai.ai_classifier._fulltext_similar_documents",
)
get_taxonomy_context(document, user=None)
mock_retrieve.assert_called_once()
mock_fulltext.assert_not_called()
class TestGetTaxonomyContextVisibility:
"""get_taxonomy_context must not materialize every visible document id
for a user who can already see the whole library: a superuser (like no
@@ -379,6 +445,7 @@ class TestGetTaxonomyContextVisibility:
"""
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_skips_permission_lookup_for_superuser(
self,
mocker: pytest_mock.MockerFixture,
@@ -397,17 +464,18 @@ class TestGetTaxonomyContextVisibility:
"paperless_ai.ai_classifier.retrieve_similar_nodes",
return_value=[],
)
mock_get_objects = mocker.patch(
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
mock_permitted = mocker.patch(
"paperless_ai.ai_classifier.permitted_object_ids",
)
user = UserFactory.create(is_superuser=True)
get_taxonomy_context(document, user)
mock_get_objects.assert_not_called()
mock_permitted.assert_not_called()
assert mock_retrieve.call_args.kwargs["document_ids"] is None
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_skips_permission_lookup_when_no_user(
self,
mocker: pytest_mock.MockerFixture,
@@ -426,16 +494,17 @@ class TestGetTaxonomyContextVisibility:
"paperless_ai.ai_classifier.retrieve_similar_nodes",
return_value=[],
)
mock_get_objects = mocker.patch(
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
mock_permitted = mocker.patch(
"paperless_ai.ai_classifier.permitted_object_ids",
)
get_taxonomy_context(document, None)
mock_get_objects.assert_not_called()
mock_permitted.assert_not_called()
assert mock_retrieve.call_args.kwargs["document_ids"] is None
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_restricts_to_visible_documents_for_non_superuser(
self,
mocker: pytest_mock.MockerFixture,
@@ -446,7 +515,7 @@ class TestGetTaxonomyContextVisibility:
WHEN:
- get_taxonomy_context() is called
THEN:
- The user's visible document ids are looked up and passed to
- The user's permitted document ids are looked up and passed to
retrieve_similar_nodes() as a restriction
"""
document = DocumentFactory.create(content="Some content")
@@ -454,21 +523,186 @@ class TestGetTaxonomyContextVisibility:
"paperless_ai.ai_classifier.retrieve_similar_nodes",
return_value=[],
)
mock_queryset = mocker.MagicMock()
mock_queryset.values_list.return_value = [1, 2, 3]
mock_get_objects = mocker.patch(
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
return_value=mock_queryset,
mock_permitted = mocker.patch(
"paperless_ai.ai_classifier.permitted_object_ids",
return_value=[1, 2, 3],
)
user = UserFactory.create(is_superuser=False)
get_taxonomy_context(document, user)
mock_get_objects.assert_called_once_with(user, "view_document", Document)
mock_permitted.assert_called_once_with(user, Document, "view_document")
assert mock_retrieve.call_args.kwargs["document_ids"] == [1, 2, 3]
@pytest.mark.django_db
class TestFulltextSimilarDocuments:
"""_fulltext_similar_documents is the no-embedding-backend fallback: it
asks the Tantivy full-text index for "More Like This" neighbours instead
of the vector store, and synthesizes a rank-based weight since Tantivy's
more_like_this_ids returns only an ordered id list, no scores.
"""
@pytest.fixture
def fulltext_backend(
self,
mocker: pytest_mock.MockerFixture,
) -> Generator[TantivyBackend, None, None]:
"""An in-memory Tantivy backend, wired up as the module-level
singleton _fulltext_similar_documents resolves via get_backend()."""
backend = TantivyBackend(path=None)
backend.open()
mocker.patch("documents.search.get_backend", return_value=backend)
try:
yield backend
finally:
backend.close()
def test_ranks_by_rank_based_weight_descending(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A source document and two similar documents indexed in Tantivy
WHEN:
- _fulltext_similar_documents() is called
THEN:
- Each result's weight reflects its rank (first result weighted
higher than the second), not a raw similarity score
"""
source = DocumentFactory.create(content="quarterly financial report details")
first = DocumentFactory.create(content="quarterly financial report details")
second = DocumentFactory.create(content="financial report")
for doc in (source, first, second):
fulltext_backend.add_or_update(doc)
result = _fulltext_similar_documents(source, user=None, top_k=5)
assert len(result) == 2
weight_by_id = {s["document_id"]: s["weight"] for s in result}
assert weight_by_id[first.pk] > weight_by_id[second.pk]
def test_excludes_source_document(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A source document indexed in Tantivy with no other documents
WHEN:
- _fulltext_similar_documents() is called
THEN:
- An empty list is returned - the source document is never its
own similar document
"""
source = DocumentFactory.create(content="unique unrelated content")
fulltext_backend.add_or_update(source)
result = _fulltext_similar_documents(source, user=None, top_k=5)
assert result == []
def test_empty_index_returns_empty_list(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A document that has never been indexed (fresh/empty Tantivy index)
WHEN:
- _fulltext_similar_documents() is called
THEN:
- An empty list is returned rather than raising
"""
source = DocumentFactory.create(content="never indexed")
result = _fulltext_similar_documents(source, user=None, top_k=5)
assert result == []
def test_respects_top_k_limit(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A source document and four similar documents indexed
WHEN:
- _fulltext_similar_documents() is called with top_k=2
THEN:
- At most 2 results are returned
"""
source = DocumentFactory.create(content="shared overlapping keyword text")
fulltext_backend.add_or_update(source)
for _ in range(4):
fulltext_backend.add_or_update(
DocumentFactory.create(content="shared overlapping keyword text"),
)
result = _fulltext_similar_documents(source, user=None, top_k=2)
assert len(result) == 2
def test_result_shape_is_similar_document(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A source document and one similar document indexed
WHEN:
- _fulltext_similar_documents() is called
THEN:
- Each result is a SimilarDocument (document_id + weight only)
"""
source = DocumentFactory.create(content="shared content phrase")
other = DocumentFactory.create(content="shared content phrase")
fulltext_backend.add_or_update(source)
fulltext_backend.add_or_update(other)
result = _fulltext_similar_documents(source, user=None, top_k=5)
# rank 0 (the only/best result) with top_k=5 -> weight = top_k - rank = 5.0,
# per the "first result gets top_k, the last gets 1" formula.
assert result == [SimilarDocument(document_id=other.pk, weight=5.0)]
def test_superuser_sees_other_users_documents(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A source document owned by one user and a similar document
owned by a different user, with no sharing between them
WHEN:
- _fulltext_similar_documents() is called with a superuser
THEN:
- The other user's document is still returned as a similar
document - a superuser must not be narrowed by the backend's
owner-based permission filter
"""
owner = UserFactory.create()
other_owner = UserFactory.create()
superuser = UserFactory.create(is_superuser=True)
source = DocumentFactory.create(
content="shared content phrase",
owner=owner,
)
other = DocumentFactory.create(
content="shared content phrase",
owner=other_owner,
)
fulltext_backend.add_or_update(source)
fulltext_backend.add_or_update(other)
result = _fulltext_similar_documents(source, user=superuser, top_k=5)
assert [s["document_id"] for s in result] == [other.pk]
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
@patch("paperless_ai.ai_classifier.retrieve_similar_nodes")
def test_get_taxonomy_context_retrieval_failure_degrades_to_no_hints(mock_retrieve):
"""
@@ -495,6 +729,7 @@ def test_get_taxonomy_context_retrieval_failure_degrades_to_no_hints(mock_retrie
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
@patch("paperless_ai.ai_classifier.build_taxonomy_candidates")
@patch("paperless_ai.ai_classifier.retrieve_similar_nodes")
def test_get_taxonomy_context_candidate_building_failure_degrades_to_no_hints(
+33 -31
View File
@@ -1,5 +1,4 @@
import json
from types import SimpleNamespace
import pytest
import pytest_mock
@@ -11,6 +10,7 @@ from documents.tests.factories import StoragePathFactory
from documents.tests.factories import TagFactory
from documents.tests.factories import UserFactory
from paperless_ai.taxonomy import AssignedMetadata
from paperless_ai.taxonomy import SimilarDocument
from paperless_ai.taxonomy import TaxonomyCandidates
from paperless_ai.taxonomy import build_taxonomy_candidates
from paperless_ai.taxonomy import format_taxonomy_for_prompt
@@ -132,9 +132,8 @@ class TestGetAssignedMetadata:
assert result["tags"] == ["Owned By Someone Else"]
def make_node(document_id: int, score: float) -> SimpleNamespace:
"""A stand-in for NodeWithScore: only ``.metadata``/``.score`` are read."""
return SimpleNamespace(metadata={"document_id": str(document_id)}, score=score)
def make_similar(document_id: int, weight: float) -> SimilarDocument:
return SimilarDocument(document_id=document_id, weight=weight)
@pytest.mark.django_db
@@ -170,9 +169,9 @@ class TestBuildTaxonomyCandidates:
doc_a.tags.add(tag)
doc_b = DocumentFactory.create()
doc_b.tags.add(tag)
nodes = [make_node(doc_a.pk, 0.9), make_node(doc_b.pk, 0.4)]
similar_documents = [make_similar(doc_a.pk, 0.9), make_similar(doc_b.pk, 0.4)]
result = build_taxonomy_candidates(nodes, user=None)
result = build_taxonomy_candidates(similar_documents, user=None)
assert len(result["tags"]) == 1
assert result["tags"][0]["id"] == tag.pk
@@ -197,9 +196,9 @@ class TestBuildTaxonomyCandidates:
document.tags.add(tag)
tag.name = "New Name"
tag.save()
nodes = [make_node(document.pk, 0.5)]
similar_documents = [make_similar(document.pk, 0.5)]
result = build_taxonomy_candidates(nodes, user=None)
result = build_taxonomy_candidates(similar_documents, user=None)
assert result["tags"][0]["name"] == "New Name"
@@ -219,9 +218,9 @@ class TestBuildTaxonomyCandidates:
document = DocumentFactory.create()
document.tags.add(tag)
tag.delete()
nodes = [make_node(document.pk, 0.5)]
similar_documents = [make_similar(document.pk, 0.5)]
result = build_taxonomy_candidates(nodes, user=None)
result = build_taxonomy_candidates(similar_documents, user=None)
assert result["tags"] == []
@@ -240,9 +239,12 @@ class TestBuildTaxonomyCandidates:
strong_doc.tags.add(strong_tag)
weak_doc = DocumentFactory.create()
weak_doc.tags.add(weak_tag)
nodes = [make_node(strong_doc.pk, 0.9), make_node(weak_doc.pk, 0.1)]
similar_documents = [
make_similar(strong_doc.pk, 0.9),
make_similar(weak_doc.pk, 0.1),
]
result = build_taxonomy_candidates(nodes, user=None)
result = build_taxonomy_candidates(similar_documents, user=None)
assert [c["name"] for c in result["tags"]] == ["Strong", "Weak"]
@@ -258,9 +260,9 @@ class TestBuildTaxonomyCandidates:
document = DocumentFactory.create()
for i in range(15):
document.tags.add(TagFactory.create(name=f"Tag{i}"))
nodes = [make_node(document.pk, 0.5)]
similar_documents = [make_similar(document.pk, 0.5)]
result = build_taxonomy_candidates(nodes, user=None)
result = build_taxonomy_candidates(similar_documents, user=None)
assert len(result["tags"]) == 10
@@ -274,12 +276,12 @@ class TestBuildTaxonomyCandidates:
- Only 5 correspondents are returned
"""
correspondents = CorrespondentFactory.create_batch(7)
nodes = [
make_node(DocumentFactory.create(correspondent=c).pk, 0.5)
similar_documents = [
make_similar(DocumentFactory.create(correspondent=c).pk, 0.5)
for c in correspondents
]
result = build_taxonomy_candidates(nodes, user=None)
result = build_taxonomy_candidates(similar_documents, user=None)
assert len(result["correspondents"]) == 5
@@ -294,9 +296,9 @@ class TestBuildTaxonomyCandidates:
"""
document_type = DocumentTypeFactory.create(name="Invoice")
document = DocumentFactory.create(document_type=document_type)
nodes = [make_node(document.pk, 0.5)]
similar_documents = [make_similar(document.pk, 0.5)]
result = build_taxonomy_candidates(nodes, user=None)
result = build_taxonomy_candidates(similar_documents, user=None)
assert len(result["document_types"]) == 1
assert result["document_types"][0]["id"] == document_type.pk
@@ -312,12 +314,12 @@ class TestBuildTaxonomyCandidates:
- Only 5 document_types are returned
"""
document_types = DocumentTypeFactory.create_batch(7)
nodes = [
make_node(DocumentFactory.create(document_type=dt).pk, 0.5)
similar_documents = [
make_similar(DocumentFactory.create(document_type=dt).pk, 0.5)
for dt in document_types
]
result = build_taxonomy_candidates(nodes, user=None)
result = build_taxonomy_candidates(similar_documents, user=None)
assert len(result["document_types"]) == 5
@@ -332,9 +334,9 @@ class TestBuildTaxonomyCandidates:
"""
storage_path = StoragePathFactory.create(name="Invoices")
document = DocumentFactory.create(storage_path=storage_path)
nodes = [make_node(document.pk, 0.5)]
similar_documents = [make_similar(document.pk, 0.5)]
result = build_taxonomy_candidates(nodes, user=None)
result = build_taxonomy_candidates(similar_documents, user=None)
assert len(result["storage_paths"]) == 1
assert result["storage_paths"][0]["id"] == storage_path.pk
@@ -350,12 +352,12 @@ class TestBuildTaxonomyCandidates:
- Only 5 storage_paths are returned
"""
storage_paths = StoragePathFactory.create_batch(7)
nodes = [
make_node(DocumentFactory.create(storage_path=sp).pk, 0.5)
similar_documents = [
make_similar(DocumentFactory.create(storage_path=sp).pk, 0.5)
for sp in storage_paths
]
result = build_taxonomy_candidates(nodes, user=None)
result = build_taxonomy_candidates(similar_documents, user=None)
assert len(result["storage_paths"]) == 5
@@ -375,14 +377,14 @@ class TestBuildTaxonomyCandidates:
tag = TagFactory.create(name="Restricted")
document = DocumentFactory.create()
document.tags.add(tag)
nodes = [make_node(document.pk, 0.5)]
similar_documents = [make_similar(document.pk, 0.5)]
user = UserFactory.create()
mocker.patch(
"documents.permissions.permitted_object_ids",
return_value=[], # user cannot see this tag
)
result = build_taxonomy_candidates(nodes, user=user)
result = build_taxonomy_candidates(similar_documents, user=user)
assert result["tags"] == []
@@ -412,10 +414,10 @@ class TestBuildTaxonomyCandidates:
tag.save()
document = DocumentFactory.create()
document.tags.add(tag)
nodes = [make_node(document.pk, 0.5)]
similar_documents = [make_similar(document.pk, 0.5)]
spy = mocker.patch("documents.permissions.permitted_object_ids")
result = build_taxonomy_candidates(nodes, user=None)
result = build_taxonomy_candidates(similar_documents, user=None)
assert result["tags"][0]["name"] == "Owned"
spy.assert_not_called()