49 lines
1.7 KiB
Python
49 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from pathlib import Path
|
|
|
|
|
|
IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp"}
|
|
DOCUMENT_SUFFIXES = {".txt", ".md", ".csv", ".json"}
|
|
MAX_IMAGE_BYTES = 10 * 1024 * 1024
|
|
MAX_DOCUMENT_BYTES = 2 * 1024 * 1024
|
|
MAX_EXTRACTED_CHARS = 30_000
|
|
|
|
|
|
def classify_upload(filename: str, mime: str, size: int) -> str:
|
|
suffix = Path(filename).suffix.lower()
|
|
if suffix in {".step", ".stp"}:
|
|
raise ValueError("STEP/STP upload is not supported in this version")
|
|
if suffix in IMAGE_SUFFIXES or mime.startswith("image/"):
|
|
if size > MAX_IMAGE_BYTES:
|
|
raise ValueError("Image upload exceeds the 10 MB limit")
|
|
return "image"
|
|
if suffix in DOCUMENT_SUFFIXES:
|
|
if size > MAX_DOCUMENT_BYTES:
|
|
raise ValueError("Document upload exceeds the 2 MB limit")
|
|
return "document"
|
|
raise ValueError("Only PNG, JPG, WEBP, TXT, MD, CSV, and JSON uploads are supported")
|
|
|
|
|
|
def extract_document_text(data: bytes) -> str:
|
|
try:
|
|
text = data.decode("utf-8")
|
|
except UnicodeDecodeError as error:
|
|
raise ValueError("Documents must be UTF-8 text") from error
|
|
return text[:MAX_EXTRACTED_CHARS]
|
|
|
|
|
|
def attachment_record(task_id: str, filename: str, mime: str, relative_path: str, data: bytes, kind: str, extracted_path: str = "") -> dict[str, object]:
|
|
return {
|
|
"id": Path(relative_path).stem,
|
|
"task_id": task_id,
|
|
"name": filename,
|
|
"kind": kind,
|
|
"path": relative_path,
|
|
"mime": mime or "application/octet-stream",
|
|
"size": len(data),
|
|
"sha256": hashlib.sha256(data).hexdigest(),
|
|
"extracted_path": extracted_path,
|
|
}
|