api.py

import io
import logging
import zipfile
from pathlib import Path
from typing import Literal

from fastapi import FastAPI, File, UploadFile
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from starlette.concurrency import run_in_threadpool

from src.config import (MAX_DOC_CHARS, MAX_HISTORY_TURNS, MAX_MESSAGE_CHARS,
                        MAX_TURN_CHARS, MAX_UPLOAD_BYTES)
from src.engine import get_ipbot_response

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("ipbot")

ROOT = Path(__file__).resolve().parent
FRONTEND = ROOT / "frontend"
ASSETS = ROOT / "assets"

app = FastAPI(title="IPBot", docs_url=None, redoc_url=None)


@app.middleware("http")
async def security_headers(request, call_next):
    response = await call_next(request)
    # everything is served from this origin, so nothing needs to reach out except the map embed
    # this is also what stops injected markup from phoning home
    response.headers["Content-Security-Policy"] = (
        "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; "
        "img-src 'self' data:; connect-src 'self'; "
        "frame-src https://www.openstreetmap.org; frame-ancestors 'none'; base-uri 'none'"
    )
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["Referrer-Policy"] = "no-referrer"
    return response


class Turn(BaseModel):
    role: Literal["user", "assistant"]
    content: str = Field(default="", max_length=MAX_TURN_CHARS)


class ChatIn(BaseModel):
    message: str = Field(default="", max_length=MAX_MESSAGE_CHARS)
    history: list[Turn] = Field(default_factory=list, max_length=MAX_HISTORY_TURNS)
    doc_text: str = Field(default="", max_length=MAX_DOC_CHARS)


@app.post("/api/chat")
def chat(body: ChatIn):
    question = body.message.strip()
    if not question and not body.doc_text.strip():
        return {"answer": "Ask me anything about IPB and I'll look it up.", "sources": []}

    history = [{"role": turn.role, "content": turn.content} for turn in body.history]
    try:
        return get_ipbot_response(question, history, doc_text=body.doc_text)
    except Exception:
        log.exception("chat failed")
        return {"answer": "⚠️ Something went wrong. Please try again.", "sources": []}


TEXT_SUFFIXES = {"txt", "md", "csv"}
ACCEPTED = TEXT_SUFFIXES | {"pdf", "docx"}


def _read_pdf(data):
    from pypdf import PdfReader

    reader = PdfReader(io.BytesIO(data))
    if len(reader.pages) > 300:
        raise ValueError("too many pages")

    parts, size = [], 0
    for page in reader.pages:
        parts.append(page.extract_text() or "")
        size += len(parts[-1])
        # only the first MAX_DOC_CHARS are ever used, so there is no point walking a page tree built to keep us busy
        if size > MAX_DOC_CHARS * 4:
            break
    return "\n".join(parts)


def _read_docx(data):
    from docx import Document

    with zipfile.ZipFile(io.BytesIO(data)) as archive:
        if "word/document.xml" not in archive.namelist():
            raise ValueError("not a word document")
        if sum(item.file_size for item in archive.infolist()) > 50_000_000:
            raise ValueError("archive expands too far")

    return "\n".join(p.text for p in Document(io.BytesIO(data)).paragraphs)


def _extract(suffix, data):
    # the extension picked the parser, so check the bytes agree with it
    if suffix == "pdf":
        if not data.startswith(b"%PDF"):
            raise ValueError("not a pdf")
        return _read_pdf(data)

    if suffix == "docx":
        if not data.startswith(b"PK"):
            raise ValueError("not a docx")
        return _read_docx(data)

    text = data.decode("utf-8", errors="replace")
    if text.count("�") > len(text) * 0.1:
        raise ValueError("not readable text")
    return text


@app.post("/api/upload")
async def upload(file: UploadFile = File(...)):
    name = Path(file.filename or "document").name[:100]
    suffix = name.lower().rsplit(".", 1)[-1] if "." in name else ""
    if suffix not in ACCEPTED:
        return {"name": name, "text": "", "error": "Upload a PDF, Word or text file."}

    # read in slices and stop early, so an oversized upload is refused rather than pulled into memory in one piece and measured afterwards
    buf = bytearray()
    while chunk := await file.read(1 << 20):
        buf += chunk
        if len(buf) > MAX_UPLOAD_BYTES:
            return {"name": name, "text": "", "error": "File too large (max 5 MB)."}

    try:
        # parsing is blocking and CPU-bound, so keep it off the event loop
        text = await run_in_threadpool(_extract, suffix, bytes(buf))
    except Exception:
        log.exception("could not read %s", name)
        return {"name": name, "text": "", "error": "Could not read this file."}

    text = text.strip()
    if not text:
        return {"name": name, "text": "", "error": "No readable text in this file."}
    return {"name": name, "text": text[:MAX_DOC_CHARS]}


app.mount("/assets", StaticFiles(directory=str(ASSETS)), name="assets")
app.mount("/static", StaticFiles(directory=str(FRONTEND)), name="static")


@app.get("/")
def index():
    return FileResponse(str(FRONTEND / "index.html"))