File attachments: images for vision, PDFs and text into the prompt
Drag, paste or pick a file in the composer. Images go to vision models as multimodal content parts; PDFs and text files have their content extracted and placed in the prompt. Verified end to end against gemma4-e4b-q8 on llama-swap: given a drawing and a text file, it named the red square and blue circle and read the number out of the document. Type is decided by inspecting the bytes, never the filename or the browser's Content-Type -- a .png full of text is stored as text. Images are downscaled to 1400px and re-encoded: a phone photo is several megabytes of base64, which is slow and a large slice of the context window. PDF text is extracted once, at upload, and stored; re-extracting per request would let a reply change because a parser was upgraded. Design points worth keeping: - Images are only sent to models an administrator has marked `vision`. This is not graceful degradation -- most endpoints reject the entire request rather than ignoring an image part. A plain text turn stays a plain string for the same reason: the list form 400s on endpoints that do not implement it. - Images reach the model as base64 data URIs, not links. A local endpoint has no route back to LLeMbas, and a hosted one has no credentials for it. - Non-images are served Content-Disposition: attachment with nosniff, so an uploaded .html can never execute in this origin. Stored names are random; the uploader's name is a label and never a path. - Uploads are unbound until the message is sent, which is what lets a file be removed beforehand. claim() only takes unclaimed rows owned by the sender, so a forged id cannot pull in someone else's file. Abandoned uploads are swept at startup. - A scanned PDF says so rather than silently contributing nothing, and truncation is declared to the model in the document tag so it can admit it did not see page 400. - "Here, look at this" with no words is a legitimate turn, so a message is only empty when it carries neither text nor files. Also fixes auto-titling, which read message["content"] as a string and would have broken on the first multimodal turn. 186 tests, ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jaroslav Beneš
@@ -0,0 +1,400 @@
|
||||
"""Storing and reading uploaded attachments.
|
||||
|
||||
Three kinds of file, each handled differently on the way to the model:
|
||||
|
||||
* **Images** are downscaled and re-encoded, then sent as multimodal content
|
||||
parts. Downscaling is not cosmetic -- a phone photo is several megabytes of
|
||||
base64, which is both slow and a large slice of the context window.
|
||||
* **PDFs** have their text extracted once, at upload. Extraction is slow and a
|
||||
reply must not silently change because a parser was upgraded later.
|
||||
* **Plain text** (including source code and CSV) is decoded and stored as-is.
|
||||
|
||||
Everything an uploader supplies is treated as hostile: the type is decided by
|
||||
inspecting the bytes rather than trusting the browser, the name on disk is
|
||||
random, and both image dimensions and PDF page counts are capped so a small
|
||||
file cannot expand into an enormous amount of work.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import logging
|
||||
import secrets
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image, UnidentifiedImageError
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session as DBSession
|
||||
|
||||
from lembas.config import settings
|
||||
from lembas.db.models import KIND_DOCUMENT, KIND_IMAGE, KIND_TEXT, Attachment
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# --- Limits ------------------------------------------------------------------
|
||||
MAX_UPLOAD_BYTES = 20 * 1024 * 1024
|
||||
|
||||
# Longest edge after downscaling. Large enough for a model to read a screenshot
|
||||
# or a page of text, small enough that the base64 stays reasonable.
|
||||
MAX_IMAGE_EDGE = 1400
|
||||
JPEG_QUALITY = 85
|
||||
|
||||
# Pillow's own guard against decompression bombs: a 60,000x60,000 PNG is a few
|
||||
# KB on disk and hundreds of GB decoded.
|
||||
Image.MAX_IMAGE_PIXELS = 64_000_000
|
||||
|
||||
MAX_PDF_PAGES = 300
|
||||
# Characters of extracted text kept per document. Roughly 30k tokens, which is
|
||||
# already a large slice of most context windows; more is rarely useful and
|
||||
# frequently breaks the request outright.
|
||||
MAX_EXTRACTED_CHARS = 120_000
|
||||
|
||||
# Orphans are files uploaded into a composer that was never sent.
|
||||
ORPHAN_AGE = timedelta(hours=24)
|
||||
|
||||
IMAGE_TYPES: dict[bytes, tuple[str, str]] = {
|
||||
b"\x89PNG\r\n\x1a\n": ("image/png", ".png"),
|
||||
b"\xff\xd8\xff": ("image/jpeg", ".jpg"),
|
||||
b"GIF87a": ("image/gif", ".gif"),
|
||||
b"GIF89a": ("image/gif", ".gif"),
|
||||
}
|
||||
|
||||
# Extensions treated as text when the bytes decode cleanly as UTF-8. The list
|
||||
# exists only to pick a sensible media type; decodability is what actually
|
||||
# decides, so an unlisted extension still works.
|
||||
TEXT_EXTENSIONS = {
|
||||
".txt": "text/plain", ".md": "text/markdown", ".markdown": "text/markdown",
|
||||
".csv": "text/csv", ".tsv": "text/tab-separated-values",
|
||||
".json": "application/json", ".yaml": "text/yaml", ".yml": "text/yaml",
|
||||
".toml": "text/toml", ".ini": "text/plain", ".cfg": "text/plain",
|
||||
".xml": "text/xml", ".html": "text/plain", ".css": "text/plain",
|
||||
".py": "text/x-python", ".js": "text/javascript", ".ts": "text/typescript",
|
||||
".rs": "text/x-rust", ".go": "text/x-go", ".c": "text/x-c", ".h": "text/x-c",
|
||||
".cpp": "text/x-c++", ".java": "text/x-java", ".rb": "text/x-ruby",
|
||||
".sh": "text/x-shellscript", ".sql": "text/x-sql", ".log": "text/plain",
|
||||
}
|
||||
|
||||
|
||||
class FileError(Exception):
|
||||
"""A rejected upload, with a message fit to show the user."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class Prepared:
|
||||
"""The result of inspecting and processing an upload, before it is stored."""
|
||||
|
||||
payload: bytes
|
||||
kind: str
|
||||
media_type: str
|
||||
extension: str
|
||||
width: int = 0
|
||||
height: int = 0
|
||||
extracted_text: str = ""
|
||||
pages: int = 0
|
||||
truncated: bool = False
|
||||
extraction_error: str = ""
|
||||
|
||||
|
||||
# --- Storage -----------------------------------------------------------------
|
||||
def attachments_dir() -> Path:
|
||||
path = settings.uploads_dir / "attachments"
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def stored_path(stored_name: str) -> Path | None:
|
||||
"""Resolve a stored name to a path, refusing anything outside the directory."""
|
||||
if not stored_name or "/" in stored_name or "\\" in stored_name or stored_name.startswith("."):
|
||||
return None
|
||||
base = attachments_dir().resolve()
|
||||
path = (base / stored_name).resolve()
|
||||
try:
|
||||
path.relative_to(base)
|
||||
except ValueError:
|
||||
return None
|
||||
return path if path.is_file() else None
|
||||
|
||||
|
||||
# --- Type detection ----------------------------------------------------------
|
||||
def _detect_image(payload: bytes) -> tuple[str, str] | None:
|
||||
for signature, (media_type, extension) in IMAGE_TYPES.items():
|
||||
if payload.startswith(signature):
|
||||
return media_type, extension
|
||||
if payload[:4] == b"RIFF" and payload[8:12] == b"WEBP":
|
||||
return "image/webp", ".webp"
|
||||
return None
|
||||
|
||||
|
||||
def _looks_like_pdf(payload: bytes) -> bool:
|
||||
# The header is allowed a little leading junk by the spec, and real files
|
||||
# in the wild use it.
|
||||
return b"%PDF-" in payload[:1024]
|
||||
|
||||
|
||||
# --- Processing --------------------------------------------------------------
|
||||
def _process_image(payload: bytes) -> Prepared:
|
||||
try:
|
||||
with Image.open(io.BytesIO(payload)) as image:
|
||||
image.load()
|
||||
has_alpha = image.mode in ("RGBA", "LA", "P") and "transparency" in image.info
|
||||
# Animation is lost on re-encode; keeping only the first frame is
|
||||
# honest and is what a model would see anyway.
|
||||
frame = image.convert("RGBA" if has_alpha else "RGB")
|
||||
|
||||
width, height = frame.size
|
||||
longest = max(width, height)
|
||||
if longest > MAX_IMAGE_EDGE:
|
||||
scale = MAX_IMAGE_EDGE / longest
|
||||
frame = frame.resize(
|
||||
(max(1, int(width * scale)), max(1, int(height * scale))),
|
||||
Image.LANCZOS,
|
||||
)
|
||||
|
||||
buffer = io.BytesIO()
|
||||
if has_alpha:
|
||||
frame.save(buffer, format="PNG", optimize=True)
|
||||
media_type, extension = "image/png", ".png"
|
||||
else:
|
||||
frame.save(buffer, format="JPEG", quality=JPEG_QUALITY, optimize=True)
|
||||
media_type, extension = "image/jpeg", ".jpg"
|
||||
|
||||
return Prepared(
|
||||
payload=buffer.getvalue(),
|
||||
kind=KIND_IMAGE,
|
||||
media_type=media_type,
|
||||
extension=extension,
|
||||
width=frame.width,
|
||||
height=frame.height,
|
||||
)
|
||||
except Image.DecompressionBombError as exc:
|
||||
raise FileError("That image's dimensions are implausibly large.") from exc
|
||||
except (UnidentifiedImageError, OSError, ValueError) as exc:
|
||||
raise FileError("That image could not be read. Is it corrupt?") from exc
|
||||
|
||||
|
||||
def _process_pdf(payload: bytes) -> Prepared:
|
||||
from pypdf import PdfReader
|
||||
from pypdf.errors import PdfReadError
|
||||
|
||||
prepared = Prepared(
|
||||
payload=payload, kind=KIND_DOCUMENT, media_type="application/pdf", extension=".pdf"
|
||||
)
|
||||
|
||||
try:
|
||||
reader = PdfReader(io.BytesIO(payload))
|
||||
if reader.is_encrypted:
|
||||
# An empty password unlocks a surprising number of "encrypted" PDFs.
|
||||
try:
|
||||
reader.decrypt("")
|
||||
except Exception: # noqa: BLE001 - any failure means the same thing
|
||||
prepared.extraction_error = (
|
||||
"This PDF is password-protected, so its text could not be read."
|
||||
)
|
||||
return prepared
|
||||
|
||||
prepared.pages = len(reader.pages)
|
||||
chunks: list[str] = []
|
||||
total = 0
|
||||
|
||||
for index, page in enumerate(reader.pages[:MAX_PDF_PAGES]):
|
||||
try:
|
||||
text = page.extract_text() or ""
|
||||
except Exception as exc: # noqa: BLE001 - one bad page is not fatal
|
||||
log.debug("page %d of a PDF failed to extract: %s", index, exc)
|
||||
continue
|
||||
if not text.strip():
|
||||
continue
|
||||
chunks.append(f"[page {index + 1}]\n{text.strip()}")
|
||||
total += len(text)
|
||||
if total >= MAX_EXTRACTED_CHARS:
|
||||
prepared.truncated = True
|
||||
break
|
||||
|
||||
if prepared.pages > MAX_PDF_PAGES:
|
||||
prepared.truncated = True
|
||||
|
||||
prepared.extracted_text = "\n\n".join(chunks)[:MAX_EXTRACTED_CHARS]
|
||||
|
||||
if not prepared.extracted_text.strip():
|
||||
# Almost always a scan. Saying so beats the model silently ignoring
|
||||
# a document the user believes it can read.
|
||||
prepared.extraction_error = (
|
||||
"No text could be extracted. This looks like a scanned PDF; "
|
||||
"LLeMbas does not do OCR yet."
|
||||
)
|
||||
|
||||
except PdfReadError as exc:
|
||||
prepared.extraction_error = "This file is not a readable PDF."
|
||||
log.info("unreadable PDF: %s", exc)
|
||||
except Exception as exc: # noqa: BLE001 - never let a bad file 500 the upload
|
||||
prepared.extraction_error = "This PDF could not be read."
|
||||
log.warning("unexpected PDF failure: %s", exc)
|
||||
|
||||
return prepared
|
||||
|
||||
|
||||
def _process_text(payload: bytes, filename: str) -> Prepared:
|
||||
for encoding in ("utf-8", "utf-16", "latin-1"):
|
||||
try:
|
||||
text = payload.decode(encoding)
|
||||
break
|
||||
except (UnicodeDecodeError, LookupError):
|
||||
continue
|
||||
else:
|
||||
raise FileError("That file is not text, and is not a format LLeMbas can read.")
|
||||
|
||||
# Null bytes mean this decoded by luck (latin-1 decodes any byte) and is
|
||||
# really a binary file.
|
||||
if "\x00" in text[:4096]:
|
||||
raise FileError("That file is not text, and is not a format LLeMbas can read.")
|
||||
|
||||
truncated = len(text) > MAX_EXTRACTED_CHARS
|
||||
extension = Path(filename).suffix.lower()
|
||||
|
||||
return Prepared(
|
||||
payload=payload,
|
||||
kind=KIND_TEXT,
|
||||
media_type=TEXT_EXTENSIONS.get(extension, "text/plain"),
|
||||
extension=extension if extension in TEXT_EXTENSIONS else ".txt",
|
||||
extracted_text=text[:MAX_EXTRACTED_CHARS],
|
||||
truncated=truncated,
|
||||
)
|
||||
|
||||
|
||||
def prepare(payload: bytes, filename: str) -> Prepared:
|
||||
"""Inspect an upload, decide what it is, and process it accordingly."""
|
||||
if not payload:
|
||||
raise FileError("That file is empty.")
|
||||
if len(payload) > MAX_UPLOAD_BYTES:
|
||||
raise FileError(f"Files must be under {MAX_UPLOAD_BYTES // (1024 * 1024)} MB.")
|
||||
|
||||
if _detect_image(payload) is not None:
|
||||
return _process_image(payload)
|
||||
if _looks_like_pdf(payload):
|
||||
return _process_pdf(payload)
|
||||
return _process_text(payload, filename)
|
||||
|
||||
|
||||
# --- Public API --------------------------------------------------------------
|
||||
def safe_display_name(filename: str) -> str:
|
||||
"""A filename fit to show. Never used as a path; the stored name is random."""
|
||||
cleaned = Path(filename or "file").name.strip() or "file"
|
||||
return cleaned[:300]
|
||||
|
||||
|
||||
def store(
|
||||
db: DBSession,
|
||||
*,
|
||||
user_id: str,
|
||||
chat_id: str | None,
|
||||
payload: bytes,
|
||||
filename: str,
|
||||
) -> Attachment:
|
||||
"""Process and persist an upload. Raises FileError if it is unusable."""
|
||||
prepared = prepare(payload, filename)
|
||||
|
||||
stored_name = f"{secrets.token_hex(16)}{prepared.extension}"
|
||||
(attachments_dir() / stored_name).write_bytes(prepared.payload)
|
||||
|
||||
attachment = Attachment(
|
||||
user_id=user_id,
|
||||
chat_id=chat_id,
|
||||
filename=safe_display_name(filename),
|
||||
stored_name=stored_name,
|
||||
media_type=prepared.media_type,
|
||||
size_bytes=len(prepared.payload),
|
||||
kind=prepared.kind,
|
||||
width=prepared.width,
|
||||
height=prepared.height,
|
||||
extracted_text=prepared.extracted_text,
|
||||
pages=prepared.pages,
|
||||
truncated=prepared.truncated,
|
||||
extraction_error=prepared.extraction_error,
|
||||
)
|
||||
db.add(attachment)
|
||||
db.commit()
|
||||
log.info(
|
||||
"stored %s (%s, %d bytes) for user %s",
|
||||
attachment.filename,
|
||||
attachment.kind,
|
||||
attachment.size_bytes,
|
||||
user_id,
|
||||
)
|
||||
return attachment
|
||||
|
||||
|
||||
def delete(db: DBSession, attachment: Attachment) -> None:
|
||||
path = stored_path(attachment.stored_name)
|
||||
if path is not None:
|
||||
path.unlink(missing_ok=True)
|
||||
db.delete(attachment)
|
||||
db.commit()
|
||||
|
||||
|
||||
def claim(db: DBSession, *, ids: list[str], user_id: str, message_id: str) -> list[Attachment]:
|
||||
"""Bind pending uploads to the message that was just sent.
|
||||
|
||||
Only unclaimed attachments belonging to this user are taken, so a stray or
|
||||
forged id cannot pull someone else's file into a conversation.
|
||||
"""
|
||||
if not ids:
|
||||
return []
|
||||
|
||||
pending = list(
|
||||
db.scalars(
|
||||
select(Attachment).where(
|
||||
Attachment.id.in_(ids),
|
||||
Attachment.user_id == user_id,
|
||||
Attachment.message_id.is_(None),
|
||||
)
|
||||
)
|
||||
)
|
||||
for attachment in pending:
|
||||
attachment.message_id = message_id
|
||||
db.commit()
|
||||
return pending
|
||||
|
||||
|
||||
def sweep_orphans(db: DBSession, older_than: timedelta = ORPHAN_AGE) -> int:
|
||||
"""Delete uploads that were never attached to a message.
|
||||
|
||||
A file picked in the composer and then abandoned would otherwise sit on
|
||||
disk forever.
|
||||
"""
|
||||
cutoff = datetime.now(UTC) - older_than
|
||||
orphans = list(db.scalars(select(Attachment).where(Attachment.message_id.is_(None))))
|
||||
|
||||
removed = 0
|
||||
for attachment in orphans:
|
||||
created = attachment.created_at
|
||||
if created.tzinfo is None:
|
||||
created = created.replace(tzinfo=UTC)
|
||||
if created >= cutoff:
|
||||
continue
|
||||
path = stored_path(attachment.stored_name)
|
||||
if path is not None:
|
||||
path.unlink(missing_ok=True)
|
||||
db.delete(attachment)
|
||||
removed += 1
|
||||
|
||||
if removed:
|
||||
db.commit()
|
||||
log.info("swept %d orphaned upload(s)", removed)
|
||||
return removed
|
||||
|
||||
|
||||
def data_uri(attachment: Attachment) -> str | None:
|
||||
"""Base64 data URI for an image, as sent to a vision model.
|
||||
|
||||
A data URI rather than a link back to this server: a local endpoint has no
|
||||
route to LLeMbas, and a hosted one has no credentials for it.
|
||||
"""
|
||||
import base64
|
||||
|
||||
path = stored_path(attachment.stored_name)
|
||||
if path is None:
|
||||
return None
|
||||
encoded = base64.b64encode(path.read_bytes()).decode("ascii")
|
||||
return f"data:{attachment.media_type};base64,{encoded}"
|
||||
Reference in New Issue
Block a user