67 lines
2.3 KiB
Python
67 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from pathlib import Path
|
|
|
|
|
|
IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".svg"}
|
|
DOCUMENT_SUFFIXES = {".txt", ".md", ".csv", ".json", ".pdf", ".dxf", ".step", ".stp"}
|
|
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 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, SVG, TXT, MD, CSV, JSON, PDF, DXF, and STEP/STP uploads are supported")
|
|
|
|
|
|
def extract_document_text(data: bytes) -> str:
|
|
# PDF extraction is optional so the API can still accept a reference when
|
|
# the local PDF dependency is unavailable; the binary remains available as
|
|
# the primary attachment artifact.
|
|
if data.startswith(b"%PDF"):
|
|
try:
|
|
from pypdf import PdfReader
|
|
import io
|
|
return "\n".join(page.extract_text() or "" for page in PdfReader(io.BytesIO(data)).pages)[:MAX_EXTRACTED_CHARS]
|
|
except Exception:
|
|
return "[PDF reference uploaded; text extraction unavailable]"
|
|
try:
|
|
text = data.decode("utf-8")
|
|
except UnicodeDecodeError as error:
|
|
return "[Binary CAD reference uploaded; structured geometry inspection is deferred to the CAD reference analyzer]"
|
|
return text[:MAX_EXTRACTED_CHARS]
|
|
|
|
|
|
def attachment_record(
|
|
conversation_id: str,
|
|
filename: str,
|
|
mime: str,
|
|
relative_path: str,
|
|
data: bytes,
|
|
kind: str,
|
|
extracted_path: str = "",
|
|
metadata: dict[str, object] | None = None,
|
|
) -> dict[str, object]:
|
|
return {
|
|
"id": Path(relative_path).stem,
|
|
"conversation_id": conversation_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,
|
|
**(metadata or {}),
|
|
}
|