"""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 store_text( db: DBSession, *, user_id: str, chat_id: str | None, filename: str, text: str, truncated: bool = False, source_note: str = "", ) -> Attachment: """Attach text that did not arrive as a file -- a fetched web page. Written to disk like any other attachment so it can be downloaded and so there is one cleanup path, rather than a second kind of attachment that exists only in the database. """ body = text[:MAX_EXTRACTED_CHARS] payload = body.encode("utf-8") stored_name = f"{secrets.token_hex(16)}.txt" (attachments_dir() / stored_name).write_bytes(payload) attachment = Attachment( user_id=user_id, chat_id=chat_id, filename=safe_display_name(filename), stored_name=stored_name, media_type="text/plain", size_bytes=len(payload), kind=KIND_TEXT, # The URL leads the text so the model can cite it, and so the reader # can see where an attachment called "Some Page.txt" came from. extracted_text=f"Source: {source_note}\n\n{body}" if source_note else body, truncated=truncated, ) db.add(attachment) db.commit() return attachment def copy_document( db: DBSession, *, user_id: str, chat_id: str | None, document ) -> Attachment: """Copy a library document into a message being composed. A copy rather than a reference. History must not change under a conversation because a document was edited or deleted afterwards -- the same reason text is extracted once at upload instead of per request. The bytes are duplicated too, so deleting the document cannot leave a message pointing at nothing. """ from lembas.services.library import documents as documents_service stored_name = "" source = documents_service.stored_path(document.stored_name) if source is not None: stored_name = f"{secrets.token_hex(16)}{Path(source.name).suffix}" (attachments_dir() / stored_name).write_bytes(source.read_bytes()) attachment = Attachment( user_id=user_id, chat_id=chat_id, filename=document.filename or f"{document.title}.txt", stored_name=stored_name, media_type=document.media_type, size_bytes=document.size_bytes, kind=document.kind, width=document.width, height=document.height, extracted_text=document.extracted_text, pages=document.pages, truncated=document.truncated, extraction_error=document.extraction_error, ) db.add(attachment) db.commit() 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}"