Coverage for paperless_ai/taxonomy.py: 32%
67 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 09:07 +0000
1import json
2from collections import defaultdict
3from typing import TYPE_CHECKING
4from typing import Final
5from typing import TypedDict
7from django.contrib.auth.models import User
8from django.db.models import Model
9from django.db.models import Prefetch
11from documents.models import Correspondent
12from documents.models import Document
13from documents.models import DocumentType
14from documents.models import StoragePath
15from documents.models import Tag
16from documents.permissions import restrict_queryset_to_visible
17from paperless_ai.prompts.context import TaxonomyBlockPromptContext
18from paperless_ai.prompts.render import render_prompt
20if TYPE_CHECKING: 20 ↛ 21line 20 didn't jump to line 21 because the condition on line 20 was never true
21 from llama_index.core.schema import NodeWithScore
24MAX_TAG_CANDIDATES: Final = 10
25MAX_SINGLE_VALUE_CANDIDATES: Final = 5
28class TaxonomyCandidate(TypedDict):
29 id: int
30 name: str
31 weight: float
34class SimilarDocument(TypedDict):
35 document_id: int
36 weight: float
39class TaxonomyCandidates(TypedDict):
40 tags: list[TaxonomyCandidate]
41 document_types: list[TaxonomyCandidate]
42 correspondents: list[TaxonomyCandidate]
43 storage_paths: list[TaxonomyCandidate]
46def empty_taxonomy_candidates() -> TaxonomyCandidates:
47 """No candidates in any category - what callers use when retrieval was
48 skipped or failed."""
49 return TaxonomyCandidates(
50 tags=[],
51 document_types=[],
52 correspondents=[],
53 storage_paths=[],
54 )
57def _node_document_weights(nodes: list["NodeWithScore"]) -> list[SimilarDocument]:
58 """Sum each node's similarity score into its document_id (a document can
59 appear via multiple chunks/nodes) and return one SimilarDocument per
60 distinct document_id."""
61 weights: dict[int, float] = defaultdict(float)
62 for node in nodes:
63 document_id = node.metadata.get("document_id")
64 if document_id is None: # pragma: no cover
65 # Every node the indexing pipeline builds always sets
66 # document_id; this guards a malformed/partial vec0 row that
67 # shouldn't occur given the current schema.
68 continue
69 try:
70 weights[int(document_id)] += float(node.score or 0.0)
71 except (TypeError, ValueError): # pragma: no cover
72 continue
73 return sorted(
74 (
75 SimilarDocument(document_id=document_id, weight=weight)
76 for document_id, weight in weights.items()
77 ),
78 key=lambda similar: similar["weight"],
79 reverse=True,
80 )
83def _visible_ranked_candidates(
84 weighted_ids: dict[int, float],
85 model: type[Model],
86 perm: str,
87 user: User | None,
88 limit: int,
89) -> list[TaxonomyCandidate]:
90 """Drop anything ``user`` may not see, resolve the survivors' names, and
91 return them ranked by descending weight and capped at ``limit``.
93 The visibility check restricts the query to just this small
94 weighted_ids set rather than materializing every id `user` may see
95 installation-wide - resolving names and checking visibility is one
96 query either way, so this never pays for scanning the whole taxonomy.
97 """
98 if not weighted_ids:
99 return []
100 visible_queryset = restrict_queryset_to_visible(
101 model.objects.filter(pk__in=weighted_ids),
102 user,
103 perm,
104 )
105 id_to_name = dict(visible_queryset.values_list("id", "name"))
106 candidates = [
107 TaxonomyCandidate(id=object_id, name=id_to_name[object_id], weight=weight)
108 for object_id, weight in weighted_ids.items()
109 if object_id in id_to_name
110 ]
111 candidates.sort(key=lambda c: c["weight"], reverse=True)
112 return candidates[:limit]
115def build_taxonomy_candidates(
116 similar_documents: list[SimilarDocument],
117 user: User | None,
118) -> TaxonomyCandidates:
119 """Resolve each similar document's id to a live Document, read its
120 *current* tags/type/correspondent/storage_path via the ORM (never any
121 possibly-stale names an adapter's source might have cached), weight each
122 distinct taxonomy object by aggregate similarity weight, permission-filter
123 against what ``user`` can see, and return each category ranked by weight
124 and capped. ``similar_documents`` may come from either the vector-RAG
125 adapter or the full-text fallback adapter - both produce this same shape.
126 """
127 if not similar_documents:
128 return empty_taxonomy_candidates()
130 # Both adapters guarantee at most one SimilarDocument per document_id, so
131 # this never silently drops a duplicate's weight.
132 document_weights: dict[int, float] = {
133 s["document_id"]: s["weight"] for s in similar_documents
134 }
136 # Only .tags.all() needs prefetching (a reverse M2M, one extra query for
137 # the whole batch). document_type/correspondent/storage_path are read
138 # below via their *_id columns (neighbour.document_type_id, etc.), which
139 # are already present on each Document row with no join - so this
140 # deliberately does NOT select_related() those three; it would fetch the
141 # full related row just to reach an id already sitting on `neighbour`.
142 neighbours = Document.objects.filter(
143 pk__in=document_weights.keys(),
144 ).prefetch_related(
145 Prefetch("tags", queryset=Tag.objects.filter(is_inbox_tag=False)),
146 )
148 tag_weights: dict[int, float] = defaultdict(float)
149 document_type_weights: dict[int, float] = defaultdict(float)
150 correspondent_weights: dict[int, float] = defaultdict(float)
151 storage_path_weights: dict[int, float] = defaultdict(float)
153 for neighbour in neighbours:
154 weight = document_weights[neighbour.pk]
155 for tag in neighbour.tags.all():
156 tag_weights[tag.pk] += weight
157 if neighbour.document_type_id:
158 document_type_weights[neighbour.document_type_id] += weight
159 if neighbour.correspondent_id:
160 correspondent_weights[neighbour.correspondent_id] += weight
161 if neighbour.storage_path_id:
162 storage_path_weights[neighbour.storage_path_id] += weight
164 return TaxonomyCandidates(
165 tags=_visible_ranked_candidates(
166 tag_weights,
167 Tag,
168 "view_tag",
169 user,
170 MAX_TAG_CANDIDATES,
171 ),
172 document_types=_visible_ranked_candidates(
173 document_type_weights,
174 DocumentType,
175 "view_documenttype",
176 user,
177 MAX_SINGLE_VALUE_CANDIDATES,
178 ),
179 correspondents=_visible_ranked_candidates(
180 correspondent_weights,
181 Correspondent,
182 "view_correspondent",
183 user,
184 MAX_SINGLE_VALUE_CANDIDATES,
185 ),
186 storage_paths=_visible_ranked_candidates(
187 storage_path_weights,
188 StoragePath,
189 "view_storagepath",
190 user,
191 MAX_SINGLE_VALUE_CANDIDATES,
192 ),
193 )
196def format_taxonomy_for_prompt(
197 candidates: TaxonomyCandidates,
198) -> str:
199 """Render ranked candidates as a labelled prompt block.
201 Candidate names are untrusted, user-controlled data, so they are
202 JSON-serialized (id/name only - weight is an internal ranking detail)
203 rather than bullet-rendered, matching the untrusted-data handling already
204 used for document content elsewhere in this module. Returns "" when there
205 are no candidates, so callers can treat the result the same as no hints at all.
206 """
207 candidate_payload = {
208 key: [{"id": c["id"], "name": c["name"]} for c in values]
209 for key, values in candidates.items()
210 if values
211 }
213 return render_prompt(
214 TaxonomyBlockPromptContext(
215 candidate_payload_json=(
216 json.dumps(candidate_payload, ensure_ascii=False)
217 if candidate_payload
218 else ""
219 ),
220 ),
221 )