mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-08-11 05:13:18 +00:00
Ok the ai workflow module, actually does the thing
This commit is contained in:
shamoon
@@ -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
|
||||
Reference in New Issue
Block a user