Coverage for paperless_ai/base_model.py: 64%
52 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 09:07 +0000
1from typing import Any
2from typing import Final
3from typing import TypedDict
5from pydantic import BaseModel
6from pydantic import Field
7from pydantic import ValidationInfo
8from pydantic import field_validator
9from pydantic.fields import FieldInfo
11# taxonomy.py MAX_TAG_CANDIDATES = 10, prompt is "up to 3 relevant dates"
12MAX_EXISTING_IDS: Final = 10
13MAX_NEW_NAMES: Final = 8
14MAX_SINGLE_VALUE_NAMES: Final = 4
15MAX_DATES: Final = 3
16# Matches documents.models.Document.title's CharField(max_length=128).
17MAX_TITLE_LENGTH: Final = 128
20def _truncate_to_field_limit(value: Any, field: FieldInfo) -> Any:
21 """
22 Clip down to its it's declared maximum. Run as a `mode="before"` validator.
23 """
24 limit = next(
25 (m.max_length for m in field.metadata if hasattr(m, "max_length")),
26 None,
27 )
28 return (
29 value
30 if (limit is None or not isinstance(value, (list, str)))
31 else value[:limit]
32 )
35# This model is serialized into the schema handed to the LLM, so its docstring
36# and field descriptions are instructions for the model. Keep implementation
37# details in code comments instead.
38class DocumentClassifierSchema(BaseModel):
39 """Classification suggestions for a single document."""
41 title: str = Field(
42 max_length=MAX_TITLE_LENGTH,
43 description=(
44 "A short, descriptive title for this document, at most "
45 f"{MAX_TITLE_LENGTH} characters."
46 ),
47 )
48 tags: list[str] = Field(
49 default_factory=list,
50 max_length=MAX_NEW_NAMES,
51 description=(
52 "All topic labels you would suggest from the document itself, e.g. "
53 "'Insurance', 'Car', 'Warranty'. Always include every suggested "
54 "name here, even when it matches an available tag."
55 ),
56 )
57 matched_tags: list[str] = Field(
58 default_factory=list,
59 max_length=MAX_NEW_NAMES,
60 description=(
61 "Names copied exactly from tags that mean the same thing as an "
62 "available tag. Align each name by position with tag_ids."
63 ),
64 )
65 tag_ids: list[int] = Field(
66 default_factory=list,
67 max_length=MAX_EXISTING_IDS,
68 description=(
69 "Available tag IDs matching matched_tags, in the same order. "
70 "Only use IDs shown in the prompt."
71 ),
72 )
73 correspondents: list[str] = Field(
74 default_factory=list,
75 max_length=MAX_SINGLE_VALUE_NAMES,
76 description=(
77 "Who this document is from or was sent to, not every party merely "
78 "mentioned. A document has a single correspondent, so give at most "
79 f"{MAX_SINGLE_VALUE_NAMES}, best first, and prefer one name over "
80 "several names for the same organisation. Always include every "
81 "suggested name here, even when it matches an available "
82 "correspondent."
83 ),
84 )
85 matched_correspondents: list[str] = Field(
86 default_factory=list,
87 max_length=MAX_SINGLE_VALUE_NAMES,
88 description=(
89 "Names copied exactly from correspondents that identify the same "
90 "entity as an available correspondent. Align each name by position "
91 "with correspondent_ids."
92 ),
93 )
94 correspondent_ids: list[int] = Field(
95 default_factory=list,
96 max_length=MAX_SINGLE_VALUE_NAMES,
97 description=(
98 "Available correspondent IDs matching matched_correspondents, in "
99 "the same order. Only use IDs shown in the prompt."
100 ),
101 )
102 document_types: list[str] = Field(
103 default_factory=list,
104 max_length=MAX_SINGLE_VALUE_NAMES,
105 description=(
106 "What kind of document this is, e.g. 'Invoice', 'Contract', 'Bank "
107 "Statement', 'Letter'. Never use its subject or sender as a "
108 "document type. A document has a single type, so give at most "
109 f"{MAX_SINGLE_VALUE_NAMES}, best first. Always include every "
110 "suggested name here, even when it matches an available document "
111 "type."
112 ),
113 )
114 matched_document_types: list[str] = Field(
115 default_factory=list,
116 max_length=MAX_SINGLE_VALUE_NAMES,
117 description=(
118 "Names copied exactly from document_types that mean the same thing "
119 "as an available document type. Align each name by position with "
120 "document_type_ids."
121 ),
122 )
123 document_type_ids: list[int] = Field(
124 default_factory=list,
125 max_length=MAX_SINGLE_VALUE_NAMES,
126 description=(
127 "Available document type IDs matching matched_document_types, in "
128 "the same order. Only use IDs shown in the prompt."
129 ),
130 )
131 storage_paths: list[str] = Field(
132 default_factory=list,
133 max_length=MAX_SINGLE_VALUE_NAMES,
134 description=(
135 "Folder-style filing location, e.g. 'Finance/Invoices'. Leave "
136 "empty unless a filing location is clearly implied - never put "
137 "tags, document types or correspondents here. A document has a "
138 f"single storage path, so give at most {MAX_SINGLE_VALUE_NAMES}, "
139 "best first. Always include every suggested name here, even when "
140 "it matches an available storage path."
141 ),
142 )
143 matched_storage_paths: list[str] = Field(
144 default_factory=list,
145 max_length=MAX_SINGLE_VALUE_NAMES,
146 description=(
147 "Names copied exactly from storage_paths that mean the same filing "
148 "location as an available storage path. Align each name by position "
149 "with storage_path_ids."
150 ),
151 )
152 storage_path_ids: list[int] = Field(
153 default_factory=list,
154 max_length=MAX_SINGLE_VALUE_NAMES,
155 description=(
156 "Available storage path IDs matching matched_storage_paths, in the "
157 "same order. Only use IDs shown in the prompt."
158 ),
159 )
160 dates: list[str] = Field(
161 default_factory=list,
162 max_length=MAX_DATES,
163 description=(
164 f"Up to {MAX_DATES} dates relevant to this document, each "
165 "formatted YYYY-MM-DD. The most important is the date the "
166 "document was issued."
167 ),
168 )
170 @field_validator(
171 "title",
172 "tags",
173 "matched_tags",
174 "tag_ids",
175 "correspondents",
176 "matched_correspondents",
177 "correspondent_ids",
178 "document_types",
179 "matched_document_types",
180 "document_type_ids",
181 "storage_paths",
182 "matched_storage_paths",
183 "storage_path_ids",
184 "dates",
185 mode="before",
186 )
187 @classmethod
188 def _truncate(cls, value: Any, info: ValidationInfo) -> Any:
189 return _truncate_to_field_limit(value, cls.model_fields[info.field_name])
192class TaxonomyChoiceDict(TypedDict):
193 """Internal representation of names and existing IDs for one taxonomy."""
195 existing_ids: list[int]
196 new_names: list[str]
199class ClassificationSuggestions(TypedDict):
200 """Internal shape used after the flat LLM response is validated."""
202 title: str
203 tags: TaxonomyChoiceDict
204 correspondents: TaxonomyChoiceDict
205 document_types: TaxonomyChoiceDict
206 storage_paths: TaxonomyChoiceDict
207 dates: list[str]
210def model_to_classification_suggestions(
211 model: DocumentClassifierSchema,
212 allowed_candidate_ids: dict[str, set[int]] | None = None,
213) -> ClassificationSuggestions:
214 """Validate optional candidate mappings and convert to the internal shape.
216 A mapping is accepted only when its name is copied from the model's own
217 complete suggestion list and its ID was actually shown for that category.
218 Invalid or unpaired mappings leave the original name untouched.
219 """
220 allowed_candidate_ids = allowed_candidate_ids or {}
222 def _choice(
223 names: list[str],
224 matched_names: list[str],
225 ids: list[int],
226 category: str,
227 ) -> TaxonomyChoiceDict:
228 remaining_names = [name for name in names if name.strip()]
229 existing_ids: list[int] = []
230 allowed_ids = allowed_candidate_ids.get(category, set())
231 for name, object_id in zip(matched_names, ids, strict=False):
232 if (
233 not name.strip()
234 or name not in remaining_names
235 or object_id not in allowed_ids
236 or object_id in existing_ids
237 ):
238 continue
239 remaining_names.remove(name)
240 existing_ids.append(object_id)
241 return TaxonomyChoiceDict(
242 existing_ids=existing_ids,
243 new_names=remaining_names,
244 )
246 return ClassificationSuggestions(
247 title=model.title,
248 tags=_choice(
249 model.tags,
250 model.matched_tags,
251 model.tag_ids,
252 "tags",
253 ),
254 correspondents=_choice(
255 model.correspondents,
256 model.matched_correspondents,
257 model.correspondent_ids,
258 "correspondents",
259 ),
260 document_types=_choice(
261 model.document_types,
262 model.matched_document_types,
263 model.document_type_ids,
264 "document_types",
265 ),
266 storage_paths=_choice(
267 model.storage_paths,
268 model.matched_storage_paths,
269 model.storage_path_ids,
270 "storage_paths",
271 ),
272 dates=model.dates,
273 )
276def classification_suggestions_to_model(
277 suggestions: ClassificationSuggestions,
278) -> DocumentClassifierSchema:
279 """Convert internal suggestions to the flat shape used for localization."""
280 return DocumentClassifierSchema(
281 title=suggestions["title"],
282 tags=suggestions["tags"]["new_names"],
283 matched_tags=[],
284 tag_ids=[],
285 correspondents=suggestions["correspondents"]["new_names"],
286 matched_correspondents=[],
287 correspondent_ids=[],
288 document_types=suggestions["document_types"]["new_names"],
289 matched_document_types=[],
290 document_type_ids=[],
291 storage_paths=suggestions["storage_paths"]["new_names"],
292 matched_storage_paths=[],
293 storage_path_ids=[],
294 dates=suggestions["dates"],
295 )