theo-agent-dashboard

attachments.py (5152B)

---

 """Attachment routes: upload, download, and data URL generation."""
 
 import base64
 import mimetypes
 import time
 import uuid
 from pathlib import Path
 
 from fastapi import APIRouter, File, Form, HTTPException, UploadFile
 from fastapi.responses import FileResponse, JSONResponse
 
 from app.config import get_config
 from app.db import Database
 from app.models import AttachmentResponse
 
 router = APIRouter(prefix="/api/attachments", tags=["attachments"])
 
 _db: Database | None = None
 
 
 def set_db(db: Database) -> None:  # pragma: no mutate: block
     """Set the module-level database instance."""
     global _db  # noqa: PLW0603
     _db = db
 
 
 _MAX_SIZE_BYTES = 50 * 1024 * 1024  # 50 MB
 _IMAGE_PREFIXES = ("image/",)
 
 
 def _ensure_db() -> Database:  # pragma: no mutate: block
     """Return the module-level database instance."""
     if _db is None:
         raise RuntimeError("Attachment DB not initialized")
     return _db
 
 
 def _attachment_dir(session_id: str) -> Path:  # pragma: no mutate: block
     """Return (and create) the storage directory for a session."""
     base = Path(get_config().data_dir) / "attachments" / session_id
     base.mkdir(parents=True, exist_ok=True)
     return base
 
 
 def _save_file(  # pragma: no mutate: block
     content: bytes,
     filename: str,
     session_id: str,
 ) -> tuple[str, str]:
     """Write file to disk and return (attachment_id, disk_path)."""
     ext = Path(filename).suffix or ""
     attachment_id = str(uuid.uuid4())
     disk_path = _attachment_dir(session_id) / f"{attachment_id}{ext}"
     disk_path.write_bytes(content)
     return attachment_id, str(disk_path)
 
 
 def _store_metadata(  # pragma: no mutate: block
     attachment_id: str,
     session_id: str,
     filename: str,
     content_type: str,
     size_bytes: int,
     disk_path: str,
     created_at: int,
 ) -> None:  # pragma: no mutate: block
     """Persist attachment metadata in SQLite."""
     db = _ensure_db()
     db.connection.execute(
         """INSERT INTO attachments
            (id, session_id, filename, content_type, size_bytes, disk_path, created_at)
            VALUES (?, ?, ?, ?, ?, ?, ?)""",
         (attachment_id, session_id, filename, content_type, size_bytes, disk_path, created_at),
     )
     db.connection.commit()
 
 
 def _fetch_metadata(attachment_id: str) -> dict | None:  # pragma: no mutate: block
     """Look up attachment metadata by id."""
     db = _ensure_db()
     rows = db.connection.execute(
         "SELECT id, filename, content_type, size_bytes, disk_path, created_at "
         "FROM attachments WHERE id = ?",
         (attachment_id,),
     ).fetchall()
     if not rows:
         return None
     row = rows[0]
     return {
         "id": row[0],
         "filename": row[1],
         "content_type": row[2],
         "size_bytes": row[3],
         "disk_path": row[4],
         "created_at": row[5],
     }
 
 
 def _is_image(content_type: str) -> bool:  # pragma: no mutate: block
     """Check if the MIME type is an image."""
     return content_type.startswith(_IMAGE_PREFIXES)
 
 
 def _to_data_url(content: bytes, content_type: str) -> str:  # pragma: no mutate: block
     """Convert binary content to a base64 data URL."""
     encoded = base64.b64encode(content).decode("ascii")
     return f"data:{content_type};base64,{encoded}"
 
 
 @router.post("", response_model=AttachmentResponse)
 async def upload_attachment(  # pragma: no mutate: block
     file: UploadFile = File(...),
     session_id: str = Form(...),
     return_data_url: str = Form("false"),
 ) -> AttachmentResponse:  # pragma: no mutate: block
     """Upload a file attachment. Rejects files > 50 MB with 413."""
     content = await file.read()
     if len(content) > _MAX_SIZE_BYTES:
         raise HTTPException(status_code=413, detail="File exceeds 50 MB limit")
 
     content_type = file.content_type or "application/octet-stream"
     created_at = int(time.time())
 
     attachment_id, disk_path = _save_file(content, file.filename, session_id)
     _store_metadata(
         attachment_id, session_id, file.filename,
         content_type, len(content), disk_path, created_at,
     )
 
     response = AttachmentResponse(
         id=attachment_id,
         filename=file.filename,
         content_type=content_type,
         size_bytes=len(content),
         created_at=created_at,
     )
 
     # Attach data URL if requested for image uploads
     if return_data_url == "true" and _is_image(content_type):
         payload = response.model_dump()
         payload["data_url"] = _to_data_url(content, content_type)
         return JSONResponse(content=payload)
 
     return response
 
 
 @router.get("/{attachment_id}")
 async def download_attachment(attachment_id: str) -> FileResponse:  # pragma: no mutate: block
     """Download an attachment by ID."""
     meta = _fetch_metadata(attachment_id)
     if meta is None:
         raise HTTPException(status_code=404, detail="Attachment not found")
 
     disk_path = Path(meta["disk_path"])
     if not disk_path.exists():
         raise HTTPException(status_code=404, detail="File missing from disk")
 
     return FileResponse(
         path=str(disk_path),
         media_type=meta["content_type"],
         filename=meta["filename"],
     )