diff --git a/src/documents/workflows/ai.py b/src/documents/workflows/ai.py new file mode 100644 index 000000000..0f9a384f6 --- /dev/null +++ b/src/documents/workflows/ai.py @@ -0,0 +1,238 @@ +import logging +from datetime import date +from datetime import datetime + +from django.contrib.auth.models import User + +from documents.models import Correspondent +from documents.models import Document +from documents.models import DocumentType +from documents.models import StoragePath +from documents.models import Tag +from documents.models import WorkflowAction +from paperless.config import AIConfig +from paperless_ai.ai_classifier import get_ai_document_classification +from paperless_ai.ai_classifier import get_llm_output_language +from paperless_ai.matching import extract_unmatched_names +from paperless_ai.matching import match_correspondents_by_name +from paperless_ai.matching import match_document_types_by_name +from paperless_ai.matching import match_storage_paths_by_name +from paperless_ai.matching import match_tags_by_name + +logger = logging.getLogger("paperless.workflows.ai") + +AISuggestionField = WorkflowAction.AISuggestionField + +# Tags use m2m relation instead +DIRECT_FIELDS: dict[str, str] = { + AISuggestionField.TITLE: "title", + AISuggestionField.CORRESPONDENT: "correspondent", + AISuggestionField.DOCUMENT_TYPE: "document_type", + AISuggestionField.STORAGE_PATH: "storage_path", + AISuggestionField.CREATED: "created", +} + + +def resolve_date(dates: list[str]) -> date | None: + """ + First usable date out of the suggestions, which are expected as + YYYY-MM-DD. Document.created is a DateField, so only one can be applied. + """ + for value in dates: + try: + return datetime.strptime(value, "%Y-%m-%d").date() + except (TypeError, ValueError): + logger.debug("Ignoring unparsable suggested date %s", value) + return None + + +def resolve_object( + model, + names: list[str], + matched: list, + *, + create_missing: bool, + owner: User | None, +): + """ + Single object from a suggestion list. The best match if there was one, else + optionally a newly-created object. StoragePaths are excluded. + """ + if matched: + return matched[0] + + if not create_missing or model is StoragePath: + return None + + unmatched = extract_unmatched_names(names, matched) + if not unmatched: + return None + + # (name, owner) is what MatchingModel is unique on + obj, created = model.objects.get_or_create( + name=unmatched[0][:128], + owner=owner, + ) + if created: + logger.info("Created %s '%s' from AI suggestion", model.__name__, obj.name) + return obj + + +def resolve_tags( + names: list[str], + matched: list[Tag], + *, + create_missing: bool, + owner: User | None, +) -> list[Tag]: + """ + Matched tags, plus newly created ones if create_missing is set. + """ + tags = list(matched) + if not create_missing: + return tags + + for name in extract_unmatched_names(names, matched): + tag, created = Tag.objects.get_or_create( + name=name[:128], + owner=owner, + ) + if created: + logger.info("Created tag '%s' from AI suggestion", tag.name) + tags.append(tag) + return tags + + +def apply_ai_suggestions_to_document( + action: WorkflowAction, + document: Document, + logging_group=None, +) -> list[str]: + """ + Get suggestions about `document` and write the chosen fields. + + Returns the names of the fields that were actually changed. + """ + selected = set(action.ai_suggestion_fields or []) + if not selected: + logger.warning( + "Workflow action %s has no AI suggestion fields selected, skipping", + action.pk, + extra={"group": logging_group}, + ) + return [] + + ai_config = AIConfig() + if not ai_config.ai_enabled: + logger.error( + "AI is not enabled, cannot apply AI suggestions for document %s", + document.pk, + extra={"group": logging_group}, + ) + return [] + + # Workflows run without a user, so we use the document owner + owner = document.owner + + try: + suggestions = get_ai_document_classification( + document, + owner, + get_llm_output_language(ai_config, owner), + ) + except Exception: + logger.exception( + "Error getting AI suggestions for document %s", + document.pk, + extra={"group": logging_group}, + ) + return [] + + overwrite = action.ai_overwrite_existing + create_missing = action.ai_create_missing + updated_fields: list[str] = [] + + def should_set(field: str) -> bool: + # The field is selected and (overwrite or it's empty) + return field in selected and ( + overwrite or getattr(document, DIRECT_FIELDS[field]) in (None, "") + ) + + if should_set(AISuggestionField.TITLE): + title = (suggestions.get("title") or "").strip() + if title: + # title is capped at 128 characters + document.title = title[:128] + updated_fields.append("title") + + if should_set(AISuggestionField.CORRESPONDENT): + names = suggestions.get("correspondents", []) + correspondent = resolve_object( + Correspondent, + names, + match_correspondents_by_name(names, owner), + create_missing=create_missing, + owner=owner, + ) + if correspondent: + document.correspondent = correspondent + updated_fields.append("correspondent") + + if should_set(AISuggestionField.DOCUMENT_TYPE): + names = suggestions.get("document_types", []) + document_type = resolve_object( + DocumentType, + names, + match_document_types_by_name(names, owner), + create_missing=create_missing, + owner=owner, + ) + if document_type: + document.document_type = document_type + updated_fields.append("document_type") + + if should_set(AISuggestionField.STORAGE_PATH): + names = suggestions.get("storage_paths", []) + storage_path = resolve_object( + StoragePath, + names, + match_storage_paths_by_name(names, owner), + create_missing=create_missing, + owner=owner, + ) + if storage_path: + document.storage_path = storage_path + updated_fields.append("storage_path") + + if should_set(AISuggestionField.CREATED): + created = resolve_date(suggestions.get("dates", [])) + if created: + document.created = created + updated_fields.append("created") + + if updated_fields: + # save fields and update modified + document.save(update_fields=[*updated_fields, "modified"]) + + if AISuggestionField.TAGS in selected: + names = suggestions.get("tags", []) + tags = resolve_tags( + names, + match_tags_by_name(names, owner), + create_missing=create_missing, + owner=owner, + ) + if tags: + # Suggested tags are always added, so overwrite_existing + # does not really apply here + document.add_nested_tags(tags) + updated_fields.append("tags") + + logger.info( + "Applied AI suggestions %s to document %s", + updated_fields or "(none)", + document.pk, + extra={"group": logging_group}, + ) + + return updated_fields