Coverage for paperless_ai/ai_classifier.py: 31%
105 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 09:07 +0000
1import logging
3from django.conf import settings
4from django.contrib.auth.models import User
6from documents.models import Document
7from documents.permissions import permitted_object_ids
8from documents.permissions import restrict_queryset_to_visible
9from documents.permissions import user_is_unrestricted
10from paperless.config import AIConfig
11from paperless_ai.base_model import ClassificationSuggestions
12from paperless_ai.base_model import TaxonomyChoiceDict
13from paperless_ai.base_model import classification_suggestions_to_model
14from paperless_ai.client import AIClient
15from paperless_ai.db import db_connection_released
16from paperless_ai.indexing import retrieve_similar_nodes
17from paperless_ai.indexing import truncate_content
18from paperless_ai.prompts.context import ClassificationPromptContext
19from paperless_ai.prompts.context import LocalizationPromptContext
20from paperless_ai.prompts.context import RagContextPromptContext
21from paperless_ai.prompts.render import render_prompt
22from paperless_ai.taxonomy import SimilarDocument
23from paperless_ai.taxonomy import TaxonomyCandidates
24from paperless_ai.taxonomy import _node_document_weights
25from paperless_ai.taxonomy import build_taxonomy_candidates
26from paperless_ai.taxonomy import empty_taxonomy_candidates
27from paperless_ai.taxonomy import format_taxonomy_for_prompt
29logger = logging.getLogger("paperless_ai.rag_classifier")
31# Neighbours retrieved for taxonomy-candidate weighting, decoupled from
32# get_taxonomy_context's max_docs (which caps how many of those same
33# neighbours get their text spliced into the RAG context block). A wider
34# pool of weighted neighbours gives build_taxonomy_candidates() more signal
35# for which tags/correspondents/etc. actually cluster around this document,
36# while the ranked candidate lists it returns stay capped by
37# taxonomy.MAX_TAG_CANDIDATES / MAX_SINGLE_VALUE_CANDIDATES regardless of
38# how many neighbours went in - so raising this does not by itself grow the
39# prompt.
40TAXONOMY_CANDIDATE_TOP_K = 15
43def _fulltext_similar_documents(
44 document: Document,
45 user: User | None,
46 top_k: int,
47) -> list[SimilarDocument]:
48 """Rank-based fallback when no embedding backend is configured. Uses
49 Tantivy's "More Like This" (term-overlap similarity) instead of vector
50 similarity - cruder, but far better than no candidates at all.
51 more_like_this_ids returns only a ranked ID list, no scores, so weight is
52 synthesized from rank (descending from top_k) rather than claiming a
53 similarity magnitude that doesn't exist. An unrestricted user (none, or an
54 active superuser - see user_is_unrestricted) is normalized to ``None``
55 before calling, since the backend's permission filter has no superuser
56 short-circuit of its own. Results are re-checked with
57 restrict_queryset_to_visible() since Tantivy's indexed permission fields
58 lag the DB via async reindexing.
59 """
60 from documents.search import get_backend
62 unrestricted = user_is_unrestricted(user)
63 search_user = None if unrestricted else user
64 backend = get_backend()
65 similar_ids = backend.more_like_this_ids(
66 document.pk,
67 user=search_user,
68 limit=top_k,
69 )
70 if not unrestricted:
71 allowed_ids = set(
72 restrict_queryset_to_visible(
73 Document.objects.filter(pk__in=similar_ids),
74 user,
75 "view_document",
76 ).values_list("pk", flat=True),
77 )
78 similar_ids = [doc_id for doc_id in similar_ids if doc_id in allowed_ids]
79 return [
80 SimilarDocument(document_id=doc_id, weight=float(top_k - rank))
81 for rank, doc_id in enumerate(similar_ids)
82 ]
85def get_language_name(language_code: str) -> str:
86 normalized_language_code = language_code.lower()
87 for code, name in settings.LANGUAGES:
88 if code.lower() == normalized_language_code:
89 return str(name)
90 return language_code
93def get_llm_output_language(ai_config: AIConfig, user: User | None) -> str | None:
94 """
95 Language to localize LLM output into: the configured language, falling back
96 to the user's own UI language when unset.
97 """
98 output_language = ai_config.llm_output_language
99 if (
100 not output_language
101 and user is not None
102 and hasattr(user, "ui_settings")
103 and isinstance(user.ui_settings.settings, dict)
104 ):
105 output_language = user.ui_settings.settings.get("language")
106 return output_language
109def build_prompt_without_rag(
110 document: Document,
111 config: AIConfig,
112 candidates: TaxonomyCandidates | None = None,
113) -> str:
114 filename = document.filename or ""
115 content = truncate_content(
116 document.content[:4000] or "",
117 chunk_size=config.llm_embedding_chunk_size,
118 context_size=config.llm_context_size,
119 )
121 taxonomy_block = (
122 format_taxonomy_for_prompt(candidates) if candidates is not None else ""
123 )
124 has_candidates = candidates is not None and any(candidates.values())
126 return render_prompt(
127 ClassificationPromptContext(
128 filename=filename,
129 content=content,
130 taxonomy_block=taxonomy_block,
131 has_candidates=has_candidates,
132 ),
133 )
136def build_prompt_with_rag(
137 document: Document,
138 config: AIConfig,
139 candidates: TaxonomyCandidates | None = None,
140 context: str = "",
141) -> str:
142 base_prompt = build_prompt_without_rag(
143 document,
144 config,
145 candidates=candidates,
146 )
147 truncated_context = truncate_content(
148 context,
149 chunk_size=config.llm_embedding_chunk_size,
150 context_size=config.llm_context_size,
151 )
153 return render_prompt(
154 RagContextPromptContext(
155 base_prompt=base_prompt,
156 context=truncated_context,
157 ),
158 )
161def build_localization_prompt(
162 suggestions: ClassificationSuggestions,
163 output_language: str,
164) -> str:
165 """Render internal suggestions in the same flat shape the model returns.
166 Only the name fields and title are used from the localized response; the
167 merge step always keeps the original ID fields.
168 """
169 language_name = get_language_name(output_language)
170 model_suggestions = classification_suggestions_to_model(suggestions)
171 return render_prompt(
172 LocalizationPromptContext(
173 language_name=language_name,
174 suggestions_json=model_suggestions.model_dump_json(),
175 ),
176 )
179def get_taxonomy_context(
180 document: Document,
181 user: User | None = None,
182 max_docs: int = 5,
183) -> tuple[TaxonomyCandidates, str]:
184 """One retrieval feeds both taxonomy candidates and RAG text context. Uses
185 vector similarity when an embedding backend is configured, otherwise
186 falls back to Tantivy full-text "More Like This" similarity - see
187 _fulltext_similar_documents. On any retrieval failure, degrades to empty
188 candidates/context rather than propagating the exception - neither a
189 vector-store outage nor a search-index issue should block classification,
190 only its context-assisted enrichment.
191 """
192 ai_config = AIConfig()
193 try:
194 if ai_config.llm_embedding_backend:
195 # None means "no restriction" to retrieve_similar_nodes. An
196 # unrestricted user (no user at all, or an active superuser -- see
197 # user_is_unrestricted) can see every document, so skip
198 # materializing every visible pk into a Python list and passing it
199 # through as an IN filter: for a large library that is a wasted
200 # quadratic scan in the vector store at best, and past ~32,763
201 # documents a hard sqlite3.OperationalError (SQLite's
202 # bound-parameter limit) at worst.
203 # permitted_object_ids() has its own superuser shortcut that would
204 # return every Document's id anyway, so this changes nothing about
205 # which documents are considered -- only how we get there.
206 visible_document_ids = (
207 None
208 if user_is_unrestricted(user)
209 else list(permitted_object_ids(user, Document, "view_document"))
210 )
211 nodes = retrieve_similar_nodes(
212 document,
213 top_k=TAXONOMY_CANDIDATE_TOP_K,
214 document_ids=visible_document_ids,
215 )
216 similar_documents = _node_document_weights(nodes)
217 else:
218 # See _fulltext_similar_documents: it applies its own permission
219 # filter via `user`, so no visible-document-id list is needed here.
220 similar_documents = _fulltext_similar_documents(
221 document,
222 user,
223 top_k=TAXONOMY_CANDIDATE_TOP_K,
224 )
226 candidates = build_taxonomy_candidates(similar_documents, user)
228 # similar_documents is already ordered by descending weight; don't lose it.
229 similar_document_ids = [s["document_id"] for s in similar_documents]
230 similar_documents_by_id = Document.objects.in_bulk(similar_document_ids)
231 similar_docs = [
232 similar_documents_by_id[document_id]
233 for document_id in similar_document_ids
234 if document_id in similar_documents_by_id
235 ][:max_docs]
236 context_blocks = []
237 for similar in similar_docs:
238 text = similar.content[:1000] or ""
239 title = similar.title or similar.filename or "Untitled"
240 context_blocks.append(f"TITLE: {title}\n{text}")
241 except Exception:
242 logger.exception(
243 "Failed to retrieve similar-document context for document %s; "
244 "continuing without taxonomy candidates or similar-document context.",
245 document.pk,
246 )
247 return empty_taxonomy_candidates(), ""
249 return candidates, "\n\n".join(context_blocks)
252def parse_ai_response(raw: dict) -> ClassificationSuggestions:
253 """``raw`` is AIClient.run_llm_query()'s validated internal-shape result.
254 This gives the rest of the module a named, typed boundary instead of
255 passing the client's bare dict straight through everywhere.
256 """
258 def _choice(value: dict | None) -> TaxonomyChoiceDict:
259 value = value or {}
260 return TaxonomyChoiceDict(
261 existing_ids=value.get("existing_ids", []),
262 new_names=value.get("new_names", []),
263 )
265 return ClassificationSuggestions(
266 title=raw.get("title", ""),
267 tags=_choice(raw.get("tags")),
268 correspondents=_choice(raw.get("correspondents")),
269 document_types=_choice(raw.get("document_types")),
270 storage_paths=_choice(raw.get("storage_paths")),
271 dates=raw.get("dates", []),
272 )
275def _candidate_id_allowlist(
276 candidates: TaxonomyCandidates,
277) -> dict[str, set[int]]:
278 """Candidate IDs grouped by category for validating model mappings."""
279 return {
280 "tags": {candidate["id"] for candidate in candidates["tags"]},
281 "document_types": {
282 candidate["id"] for candidate in candidates["document_types"]
283 },
284 "correspondents": {
285 candidate["id"] for candidate in candidates["correspondents"]
286 },
287 "storage_paths": {candidate["id"] for candidate in candidates["storage_paths"]},
288 }
291def get_ai_document_classification(
292 document: Document,
293 user: User | None = None,
294 output_language: str | None = None,
295) -> ClassificationSuggestions:
296 ai_config = AIConfig()
298 candidates, context = get_taxonomy_context(document, user)
299 prompt = build_prompt_with_rag(
300 document,
301 ai_config,
302 candidates=candidates,
303 context=context,
304 )
306 client = AIClient()
307 # Hand the pooled DB connection back while the (slow) LLM query runs so it
308 # is not pinned for the call's duration; see paperless_ai.db and #12976.
309 with db_connection_released():
310 result = client.run_llm_query(
311 prompt,
312 allowed_candidate_ids=_candidate_id_allowlist(candidates),
313 )
314 suggestions = parse_ai_response(result)
315 if output_language:
316 localized = client.run_llm_query(
317 build_localization_prompt(suggestions, output_language),
318 )
319 localized_suggestions = parse_ai_response(localized)
321 def _localized_choice(field: str) -> TaxonomyChoiceDict:
322 # existing_ids always come from the ORIGINAL suggestions -
323 # never from localized_suggestions, whatever the model echoed
324 # back there. This is the concrete fix for the bug this
325 # feature exists to close: localization must never be able to
326 # corrupt an exact taxonomy match.
327 return TaxonomyChoiceDict(
328 existing_ids=suggestions[field]["existing_ids"],
329 new_names=localized_suggestions[field]["new_names"]
330 or suggestions[field]["new_names"],
331 )
333 suggestions = ClassificationSuggestions(
334 title=localized_suggestions["title"] or suggestions["title"],
335 tags=_localized_choice("tags"),
336 correspondents=suggestions["correspondents"], # never localized
337 document_types=_localized_choice("document_types"),
338 storage_paths=_localized_choice("storage_paths"),
339 dates=suggestions["dates"],
340 )
341 return suggestions