server.py

"""
FastAPI backend for the VLPS custom website.
Serves the static frontend (web/) and a JSON API over the core logic.

Run locally:  uvicorn server:app --reload --port 7860
On HF Spaces:  CMD in Dockerfile runs uvicorn on 0.0.0.0:7860
"""
import base64
import io
import os
import secrets

from fastapi import FastAPI, UploadFile, File, Form, Request
from fastapi.responses import JSONResponse, StreamingResponse, Response
from fastapi.staticfiles import StaticFiles
from PIL import Image

import core

app = FastAPI(title="VLPS Demo")

# Optional password gate. Set APP_PASSWORD (and optionally APP_USER) as a Space
# secret to require a login on a public Space. Left unset -> open (local dev).
APP_USER = os.getenv("APP_USER", "grader")
APP_PASSWORD = os.getenv("APP_PASSWORD", "")


@app.middleware("http")
async def password_gate(request: Request, call_next):
    if APP_PASSWORD:
        header = request.headers.get("authorization", "")
        ok = False
        if header.startswith("Basic "):
            try:
                user, _, pwd = base64.b64decode(header[6:]).decode().partition(":")
                ok = (secrets.compare_digest(user, APP_USER)
                      and secrets.compare_digest(pwd, APP_PASSWORD))
            except Exception:  # noqa: BLE001
                ok = False
        if not ok:
            return Response(status_code=401, headers={"WWW-Authenticate": 'Basic realm="VLPS"'})
    response = await call_next(request)
    # Always revalidate static assets so UI updates show without a hard refresh.
    response.headers["Cache-Control"] = "no-cache"
    return response


def _pil(upload: UploadFile) -> Image.Image:
    return Image.open(io.BytesIO(upload.file.read())).convert("RGB")


def _ok(data):
    return JSONResponse(data)


def _err(e):
    return JSONResponse({"error": str(e)}, status_code=200)


@app.get("/api/config")
def config():
    return {"model": core.MODEL, "user": core.USER}


@app.post("/api/search")
async def api_search(payload: dict):
    try:
        return _ok({"results": core.search_images(
            payload.get("query", ""),
            payload.get("mode", "knn"),
            int(payload.get("top_k", 6)),
        )})
    except Exception as e:  # noqa: BLE001
        return _err(e)


@app.post("/api/vqa")
async def api_vqa(image: UploadFile = File(...), question: str = Form(...),
                  system_prompt: str = Form("")):
    try:
        return _ok({"answer": core.visual_qa(_pil(image), question, system_prompt)})
    except Exception as e:  # noqa: BLE001
        return _err(e)


@app.post("/api/ocr_qa")
async def api_ocr_qa(image: UploadFile = File(...), question: str = Form(...),
                     lang: str = Form("en")):
    try:
        return _ok(core.ocr_qa(_pil(image), question, lang))
    except Exception as e:  # noqa: BLE001
        return _err(e)


@app.get("/api/dataset")
async def api_dataset(dataset: str = "textvqa", limit: int = 12):
    try:
        return _ok({"items": core.dataset_browse(dataset, limit)})
    except Exception as e:  # noqa: BLE001
        return _err(e)


@app.post("/api/dataset/search")
async def api_dataset_search(payload: dict):
    try:
        return _ok({"items": core.dataset_search(
            payload.get("dataset", "textvqa"),
            payload.get("query", ""),
            int(payload.get("k", 6)),
        )})
    except Exception as e:  # noqa: BLE001
        return _err(e)


@app.post("/api/agent")
async def api_agent(payload: dict):
    try:
        return _ok(core.agent_answer(payload.get("message", ""), payload.get("session", "")))
    except Exception as e:  # noqa: BLE001
        return _err(e)


@app.post("/api/agent/reset")
async def api_agent_reset(payload: dict):
    core.reset_agent(payload.get("session", ""))
    return _ok({"ok": True})


@app.get("/pt-image/{image_id}")
def pt_image(image_id: str):
    _, lookup = core.pt_data()
    img = lookup.get(image_id)
    if img is None:
        return JSONResponse({"error": "not found"}, status_code=404)
    buf = io.BytesIO()
    img.save(buf, format="JPEG")
    buf.seek(0)
    return StreamingResponse(buf, media_type="image/jpeg")


# Static frontend (mounted last so /api/* take precedence).
app.mount("/", StaticFiles(directory="web", html=True), name="web")