app/server.py

"""
Local web backend for the ProbabilityRAG tutor.

WHY a long-lived server: BGE-M3 + the reranker take ~40s to load. The CLI pays that on
every question; this server loads them ONCE at startup (lifespan) and holds the Qdrant
client open, so each question is fast. It just wraps the existing pipeline — no new logic.

Run: .venv/Scripts/python -m uvicorn app.server:app --port 8000
(Ollama must be running; set TORCH_DEVICE=cuda)
"""
from __future__ import annotations

import os
import sys
from contextlib import asynccontextmanager
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

# src.generate -> src.store imports torch + FlagEmbedding BEFORE qdrant_client (load order
# matters on Windows; see src/store.py). Keep this import first.
from src.generate import stream, ENABLED_MODELS, DEFAULT_MODEL  # noqa: E402
from src.store import get_model, get_reranker  # noqa: E402
from qdrant_client import QdrantClient  # noqa: E402

import json  # noqa: E402

from fastapi import Depends, FastAPI, HTTPException  # noqa: E402
from fastapi.middleware.cors import CORSMiddleware  # noqa: E402
from fastapi.responses import FileResponse, StreamingResponse  # noqa: E402
from fastapi.staticfiles import StaticFiles  # noqa: E402
from pydantic import BaseModel, Field  # noqa: E402

# Rate limiting is OPT-IN: only the public demo (PROBRAG_PUBLIC=1) pays the sqlite dependency;
# local runs get an empty dependency list and behave exactly as before.
if os.environ.get("PROBRAG_PUBLIC") == "1":
    from app.limiter import check_and_count  # noqa: E402
    _LIMIT = [Depends(check_and_count)]
else:
    _LIMIT = []

DB = ROOT / "qdrant_storage"
PDF = ROOT / "data" / "raw" / "grinstead_snell.pdf"  # served for the citation viewer
_state: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    print("[server] warming BGE-M3 + reranker ...", flush=True)
    # Loading puts weights on CPU; FlagEmbedding only moves them to the GPU on first use,
    # so without these dummy calls the FIRST question pays ~4s of model migration.
    get_model().encode(["warmup"], return_dense=True, return_sparse=True)
    get_reranker().compute_score([["warmup", "warmup"]], normalize=True, batch_size=8)
    _state["client"] = QdrantClient(path=str(DB))
    print("[server] ready.", flush=True)
    yield
    _state.clear()


app = FastAPI(title="ProbabilityRAG", lifespan=lifespan)
# Origins from env (comma-separated); default "*" keeps local dev unchanged.
_ORIGINS = [o.strip() for o in os.environ.get("PROBRAG_ORIGINS", "*").split(",") if o.strip()]
app.add_middleware(
    CORSMiddleware, allow_origins=_ORIGINS, allow_methods=["*"], allow_headers=["*"],
)


class AskRequest(BaseModel):
    question: str = Field(max_length=500)
    show_thinking: bool = False
    top_k: int = Field(default=6, ge=1, le=10)
    model: str | None = None


@app.get("/health")
def health() -> dict:
    return {"ok": "client" in _state}


def _run(req: AskRequest, style: str) -> StreamingResponse:
    # Prime the generator here so a bad model name (ValueError from _resolve, before any
    # retrieval) maps to HTTP 400 instead of surfacing mid-stream after headers are sent.
    events = stream(_state["client"], req.question.strip(), top_k=req.top_k,
                    show_thinking=req.show_thinking, style=style, model=req.model)
    try:
        first = next(events)
    except ValueError as e:
        raise HTTPException(400, str(e))
    except StopIteration:
        first = None

    def gen():
        # One SSE `data:` frame per pipeline event ({"type": "sources"|"token"|"done", ...}).
        # A generation failure (LLM API down, bad key, unknown model) becomes an "error"
        # frame instead of a silent truncation — the UI shows it and remote debugging works.
        try:
            if first is not None:
                yield f"data: {json.dumps(first)}\n\n"
            for ev in events:
                yield f"data: {json.dumps(ev)}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'message': f'{type(e).__name__}: {e}'})}\n\n"
    return StreamingResponse(gen(), media_type="text/event-stream")


@app.get("/models")
def models() -> dict:
    return {"models": ENABLED_MODELS, "default": DEFAULT_MODEL}


@app.post("/ask/stream", dependencies=_LIMIT)
def ask_stream(req: AskRequest) -> StreamingResponse:
    return _run(req, "default")


@app.post("/why/stream", dependencies=_LIMIT)
def why_stream(req: AskRequest) -> StreamingResponse:
    return _run(req, "deeper")


@app.post("/eli5/stream", dependencies=_LIMIT)
def eli5_stream(req: AskRequest) -> StreamingResponse:
    # "ELI5" follow-up: same grounded sources, plainest/fewest-words answer for a beginner.
    return _run(req, "eli5")


@app.get("/pdf")
def pdf() -> FileResponse:
    # The frontend embeds this as <iframe src="/pdf#page=N"> and jumps between cited pages.
    # content_disposition_type="inline" is required so the browser RENDERS the PDF in the
    # iframe (and honours #page=N) instead of downloading it (the FileResponse default).
    if not PDF.exists():
        raise HTTPException(404, "source PDF not found")
    return FileResponse(PDF, media_type="application/pdf", content_disposition_type="inline")


# Serve the built React bundle from the same origin (HF Spaces = one container). Mounted LAST
# so it only catches paths the API routes above didn't. Guarded on existence: in local dev the
# frontend runs on Vite (:5173) and web/dist isn't built, so this branch is simply skipped and
# nothing changes. html=True makes "/" serve index.html (SPA entry).
DIST = ROOT / "web" / "dist"
if DIST.exists():
    app.mount("/", StaticFiles(directory=str(DIST), html=True), name="static")