Coverage for documents/workflows/ai.py: 0%
125 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 09:07 +0000
1import logging
2from datetime import date
3from datetime import datetime
4from typing import TypeVar
6from django.contrib.auth.models import User
8from documents.models import Correspondent
9from documents.models import Document
10from documents.models import DocumentType
11from documents.models import MatchingModel
12from documents.models import StoragePath
13from documents.models import Tag
14from documents.models import WorkflowAction
15from paperless.config import AIConfig
16from paperless_ai.ai_classifier import get_ai_document_classification
17from paperless_ai.ai_classifier import get_llm_output_language
18from paperless_ai.matching import extract_unmatched_names
19from paperless_ai.matching import match_correspondents_by_name
20from paperless_ai.matching import match_document_types_by_name
21from paperless_ai.matching import match_storage_paths_by_name
22from paperless_ai.matching import match_tags_by_name
23from paperless_ai.matching import resolve_correspondent_ids
24from paperless_ai.matching import resolve_document_type_ids
25from paperless_ai.matching import resolve_storage_path_ids
26from paperless_ai.matching import resolve_tag_ids
28logger = logging.getLogger("paperless.workflows.ai")
30AISuggestionField = WorkflowAction.AISuggestionField
31ObjT = TypeVar("ObjT", bound=MatchingModel)
33# Tags use m2m relation instead
34DIRECT_FIELDS: dict[str, str] = {
35 AISuggestionField.TITLE: "title",
36 AISuggestionField.CORRESPONDENT: "correspondent",
37 AISuggestionField.DOCUMENT_TYPE: "document_type",
38 AISuggestionField.STORAGE_PATH: "storage_path",
39 AISuggestionField.CREATED: "created",
40}
43def resolve_date(dates: list[str]) -> date | None:
44 """
45 First usable date out of the suggestions, which are expected as
46 YYYY-MM-DD. Document.created is a DateField, so only one can be applied.
47 """
48 for value in dates:
49 try:
50 return datetime.strptime(value, "%Y-%m-%d").date()
51 except (TypeError, ValueError):
52 logger.debug("Ignoring unparsable suggested date %s", value)
53 return None
56def resolve_object(
57 model: type[ObjT],
58 names: list[str],
59 matched: list[ObjT],
60 *,
61 create_missing: bool,
62 owner: User | None,
63) -> ObjT | None:
64 """
65 Single object from a suggestion list. The best match if there was one, else
66 optionally a newly-created object. StoragePaths are excluded.
67 """
68 if matched:
69 return matched[0]
71 if not create_missing or model is StoragePath:
72 return None
74 unmatched = extract_unmatched_names(names, matched)
75 if not unmatched:
76 return None
78 # (name, owner) is what MatchingModel is unique on
79 obj, created = model.objects.get_or_create(
80 name=unmatched[0][:128],
81 owner=owner,
82 )
83 if created:
84 logger.info("Created %s '%s' from AI suggestion", model.__name__, obj.name)
85 return obj
88def resolve_tags(
89 names: list[str],
90 matched: list[Tag],
91 *,
92 create_missing: bool,
93 owner: User | None,
94) -> list[Tag]:
95 """
96 Matched tags, plus newly created ones if create_missing is set.
97 """
98 tags = list(matched)
99 if not create_missing:
100 return tags
102 for name in extract_unmatched_names(names, matched):
103 tag, created = Tag.objects.get_or_create(
104 name=name[:128],
105 owner=owner,
106 )
107 if created:
108 logger.info("Created tag '%s' from AI suggestion", tag.name)
109 tags.append(tag)
110 return tags
113def apply_ai_suggestions_to_document(
114 action: WorkflowAction,
115 document: Document,
116 logging_group=None,
117) -> list[str]:
118 """
119 Get suggestions about `document` and write the chosen fields.
121 Returns the names of the fields that were actually changed.
122 """
123 selected = set(action.ai_suggestion_fields or [])
124 if not selected:
125 logger.warning(
126 "Workflow action %s has no AI suggestion fields selected, skipping",
127 action.pk,
128 extra={"group": logging_group},
129 )
130 return []
132 ai_config = AIConfig()
133 if not ai_config.ai_enabled:
134 logger.error(
135 "AI is not enabled, cannot apply AI suggestions for document %s",
136 document.pk,
137 extra={"group": logging_group},
138 )
139 return []
141 if not document.content.strip():
142 logger.warning(
143 "Document %s has no content, skipping AI suggestions for workflow "
144 "action %s",
145 document.pk,
146 action.pk,
147 extra={"group": logging_group},
148 )
149 return []
151 # Workflows run without a user, so we use the document owner
152 owner = document.owner
154 try:
155 suggestions = get_ai_document_classification(
156 document,
157 owner,
158 get_llm_output_language(ai_config, owner),
159 )
160 except ValueError:
161 # A bad AI config will not fix itself, so swallow it rather than
162 # letting the caller retry. Timeouts, rate limits, network errors etc
163 # propagate so the queued task can back off and try again.
164 logger.exception(
165 "Invalid AI configuration, cannot get suggestions for document %s",
166 document.pk,
167 extra={"group": logging_group},
168 )
169 return []
171 overwrite = action.ai_overwrite_existing
172 create_missing = action.ai_create_missing
173 updated_fields: list[str] = []
175 def should_set(field: str) -> bool:
176 # The field is selected and (overwrite or it's empty)
177 return field in selected and (
178 overwrite or getattr(document, DIRECT_FIELDS[field]) in (None, "")
179 )
181 if should_set(AISuggestionField.TITLE):
182 title = suggestions["title"].strip()
183 if title:
184 # title is capped at 128 characters
185 document.title = title[:128]
186 updated_fields.append("title")
188 if should_set(AISuggestionField.CORRESPONDENT):
189 choice = suggestions["correspondents"]
190 names = choice["new_names"]
191 correspondent = resolve_object(
192 Correspondent,
193 names,
194 resolve_correspondent_ids(choice["existing_ids"], owner)
195 + match_correspondents_by_name(names, owner),
196 create_missing=create_missing,
197 owner=owner,
198 )
199 if correspondent:
200 document.correspondent = correspondent
201 updated_fields.append("correspondent")
203 if should_set(AISuggestionField.DOCUMENT_TYPE):
204 choice = suggestions["document_types"]
205 names = choice["new_names"]
206 document_type = resolve_object(
207 DocumentType,
208 names,
209 resolve_document_type_ids(choice["existing_ids"], owner)
210 + match_document_types_by_name(names, owner),
211 create_missing=create_missing,
212 owner=owner,
213 )
214 if document_type:
215 document.document_type = document_type
216 updated_fields.append("document_type")
218 if should_set(AISuggestionField.STORAGE_PATH):
219 choice = suggestions["storage_paths"]
220 names = choice["new_names"]
221 storage_path = resolve_object(
222 StoragePath,
223 names,
224 resolve_storage_path_ids(choice["existing_ids"], owner)
225 + match_storage_paths_by_name(names, owner),
226 create_missing=create_missing,
227 owner=owner,
228 )
229 if storage_path:
230 document.storage_path = storage_path
231 updated_fields.append("storage_path")
233 if should_set(AISuggestionField.CREATED):
234 created = resolve_date(suggestions["dates"])
235 if created:
236 document.created = created
237 updated_fields.append("created")
239 tags_to_add: list[Tag] = []
240 if AISuggestionField.TAGS in selected:
241 choice = suggestions["tags"]
242 names = choice["new_names"]
243 tags_to_add = resolve_tags(
244 names,
245 resolve_tag_ids(choice["existing_ids"], owner)
246 + match_tags_by_name(names, owner),
247 create_missing=create_missing,
248 owner=owner,
249 )
250 if tags_to_add:
251 # Suggested tags are always added, so overwrite_existing
252 # does not really apply here
253 updated_fields.append("tags")
255 if updated_fields:
256 # save fields and update modified (excluding m2m tags from update_fields)
257 direct_updated_fields = [
258 field for field in updated_fields if field in DIRECT_FIELDS.values()
259 ]
260 document.save(update_fields=[*direct_updated_fields, "modified"])
262 # Tags at the end so m2m_changed doesn't trigger db and overwrite other changes
263 if tags_to_add:
264 document.add_nested_tags(tags_to_add)
266 logger.info(
267 "Applied AI suggestions %s to document %s",
268 updated_fields or "(none)",
269 document.pk,
270 extra={"group": logging_group},
271 )
273 return updated_fields