server/app.py

"""FastAPI front end for the research agent.

A request kicks off a background job (the agent reads full PDFs, so a run takes
minutes). The browser polls for live progress and the final review. Jobs are
persisted to SQLite so completed reviews survive a container restart and
interrupted ones report a clear error.
"""

from __future__ import annotations

import json
import os
import sqlite3
import threading
import time
import uuid
from collections import defaultdict, deque
from concurrent.futures import ThreadPoolExecutor
from dataclasses import asdict, dataclass, field

from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

# --- Tunables (overridable via env) --------------------------------------- #
MAX_PAPERS_CAP = int(os.getenv("MAX_PAPERS_CAP", "12"))
MAX_DEPTH = int(os.getenv("MAX_DEPTH", "3"))
RATE_LIMIT_PER_HOUR = int(os.getenv("RATE_LIMIT_PER_HOUR", "5"))
MAX_CONCURRENT_JOBS = int(os.getenv("MAX_CONCURRENT_JOBS", "2"))
JOB_TTL_SECONDS = int(os.getenv("JOB_TTL_SECONDS", "86400"))
CACHE_TTL_SECONDS = int(os.getenv("CACHE_TTL_SECONDS", "86400"))
DB_PATH = os.getenv("JOBS_DB", "jobs.db")
# Selectable models (UI offers these); anything else falls back to the default.
ALLOWED_MODELS = [m.strip() for m in os.getenv(
    "ALLOWED_MODELS", "gemini-2.5-flash,gemini-2.5-pro"
).split(",") if m.strip()]
DEFAULT_PAPERS = int(os.getenv("DEFAULT_PAPERS", "5"))
# Cheaper model for the ~20 mechanical extraction calls (blank = use run model).
EXTRACT_MODEL = os.getenv("EXTRACT_MODEL", "").strip()

HERE = os.path.dirname(__file__)


@dataclass
class Job:
    id: str
    topic: str
    max_papers: int
    depth: int = 1
    model: str = ""
    year_min: int = 0
    style: str = "concise"
    status: str = "queued"  # queued | running | done | error
    log: list[str] = field(default_factory=list)
    review: str = ""
    gaps: list[str] = field(default_factory=list)
    papers: list[dict] = field(default_factory=list)  # bibliography metadata
    warnings: list[str] = field(default_factory=list)  # citation-check issues
    error: str = ""
    cached: bool = False
    created_at: float = field(default_factory=time.time)
    stats: dict = field(default_factory=dict)


class JobStore:
    """SQLite persistence for jobs (one JSON blob per job)."""

    def __init__(self, path: str):
        self._path = path
        self._lock = threading.Lock()
        with self._lock, self._connect() as c:
            c.execute(
                "CREATE TABLE IF NOT EXISTS jobs "
                "(id TEXT PRIMARY KEY, data TEXT, created_at REAL)"
            )
            # Any job left 'running'/'queued' from a previous process is orphaned.
            rows = c.execute("SELECT id, data FROM jobs").fetchall()
            for jid, data in rows:
                d = json.loads(data)
                if d.get("status") in ("running", "queued"):
                    d["status"] = "error"
                    d["error"] = "Interrupted by a server restart — please re-run."
                    c.execute("UPDATE jobs SET data=? WHERE id=?", (json.dumps(d), jid))

    def _connect(self):
        return sqlite3.connect(self._path, timeout=10)

    def save(self, job: "Job") -> None:
        with self._lock, self._connect() as c:
            c.execute(
                "INSERT OR REPLACE INTO jobs(id, data, created_at) VALUES (?, ?, ?)",
                (job.id, json.dumps(asdict(job)), job.created_at),
            )

    def load(self, job_id: str) -> "Job | None":
        with self._lock, self._connect() as c:
            row = c.execute("SELECT data FROM jobs WHERE id=?", (job_id,)).fetchone()
        return Job(**json.loads(row[0])) if row else None

    def find_cached(
        self, topic, max_papers, depth, model, year_min, style, ttl
    ) -> "Job | None":
        """Most recent completed job with identical params, within ``ttl`` seconds."""
        cutoff = time.time() - ttl
        with self._lock, self._connect() as c:
            rows = c.execute(
                "SELECT data FROM jobs WHERE created_at >= ? ORDER BY created_at DESC",
                (cutoff,),
            ).fetchall()
        for (data,) in rows:
            d = json.loads(data)
            if (
                d.get("status") == "done"
                and d.get("topic", "").strip().lower() == topic.strip().lower()
                and d.get("max_papers") == max_papers
                and d.get("depth") == depth
                and (d.get("model") or "") == (model or "")
                and (d.get("year_min") or 0) == year_min
                and (d.get("style") or "concise") == style
            ):
                return Job(**d)
        return None

    def prune(self, ttl: int) -> None:
        with self._lock, self._connect() as c:
            c.execute("DELETE FROM jobs WHERE created_at < ?", (time.time() - ttl,))


class JobManager:
    """Live jobs in memory (fast log appends), durably mirrored to SQLite."""

    def __init__(self) -> None:
        self._jobs: dict[str, Job] = {}
        self._lock = threading.Lock()
        self._pool = ThreadPoolExecutor(max_workers=MAX_CONCURRENT_JOBS)
        self._store = JobStore(DB_PATH)

    def submit(self, topic, max_papers, depth, model, year_min, style) -> Job:
        # Return a recent identical completed review instead of re-spending.
        cached = self._store.find_cached(
            topic, max_papers, depth, model, year_min, style, CACHE_TTL_SECONDS
        )
        if cached:
            cached.cached = True
            return cached

        job = Job(
            id=uuid.uuid4().hex[:12],
            topic=topic,
            max_papers=max_papers,
            depth=depth,
            model=model,
            year_min=year_min,
            style=style,
        )
        with self._lock:
            self._jobs[job.id] = job
        self._store.save(job)
        self._store.prune(JOB_TTL_SECONDS)
        self._pool.submit(self._run, job)
        return job

    def get(self, job_id: str) -> Job | None:
        with self._lock:
            job = self._jobs.get(job_id)
        return job or self._store.load(job_id)  # fall back to persisted survivors

    def _run(self, job: Job) -> None:
        from src.agent import build_graph
        from src.utils import cost_report, paper_meta, validate_review

        meter: dict = {}
        t0 = time.time()

        def progress(msg: str) -> None:
            job.log.append(msg)

        job.status = "running"
        self._store.save(job)
        progress(f"Starting review: {job.topic}")
        try:
            agent = build_graph(
                max_papers=job.max_papers,
                progress=progress,
                meter=meter,
                model=job.model or None,
                year_min=job.year_min,
                style=job.style,
                extract_model=EXTRACT_MODEL or None,
            )
            result = agent.invoke(
                {
                    "topic": job.topic,
                    "papers": [],
                    "review": "",
                    "gaps": [],
                    "depth": job.depth,
                    "iteration": 0,
                }
            )
            job.review = result.get("review", "")
            job.gaps = result.get("gaps", [])
            top = result.get("papers", [])[: job.max_papers]
            job.papers = [paper_meta(p) for p in top]
            if not job.review.strip():
                raise RuntimeError("No review was produced.")
            job.warnings = validate_review(job.review, top).get("issues", [])
            job.status = "done"
        except Exception as err:
            job.error = str(err)
            progress(f"Error: {err}")
            job.status = "error"
        finally:
            job.stats = {"elapsed_s": round(time.time() - t0, 1), **cost_report(meter)}
            progress(
                f"Finished in {job.stats['elapsed_s']}s · "
                f"{job.stats['llm_calls']} LLM calls · ~${job.stats['est_cost_usd']}"
            )
            self._store.save(job)


class RateLimiter:
    """Simple per-IP sliding-window limiter (in-memory, single instance)."""

    def __init__(self, limit: int, window_s: int = 3600) -> None:
        self.limit = limit
        self.window_s = window_s
        self._hits: dict[str, deque[float]] = defaultdict(deque)
        self._lock = threading.Lock()

    def allow(self, key: str) -> tuple[bool, int]:
        now = time.time()
        with self._lock:
            hits = self._hits[key]
            while hits and hits[0] < now - self.window_s:
                hits.popleft()
            if len(hits) >= self.limit:
                return False, int(self.window_s - (now - hits[0]))
            hits.append(now)
            return True, 0


jobs = JobManager()
limiter = RateLimiter(RATE_LIMIT_PER_HOUR)
app = FastAPI(title="Research Agent")


def _client_ip(request: Request) -> str:
    fwd = request.headers.get("x-forwarded-for", "")
    if fwd:
        return fwd.split(",")[0].strip()
    return request.client.host if request.client else "unknown"


@app.get("/api/healthz")
def healthz() -> dict:
    return {"ok": True}


@app.get("/", response_class=HTMLResponse)
def index() -> HTMLResponse:
    with open(os.path.join(HERE, "static", "index.html"), encoding="utf-8") as fh:
        return HTMLResponse(fh.read())


@app.post("/api/review")
async def create_review(request: Request) -> JSONResponse:
    body = await request.json()
    topic = (body.get("topic") or "").strip()
    if not topic:
        return JSONResponse({"error": "Please enter a topic."}, status_code=400)
    if len(topic) > 300:
        return JSONResponse({"error": "Topic is too long."}, status_code=400)

    try:
        max_papers = int(body.get("max_papers", DEFAULT_PAPERS))
    except (TypeError, ValueError):
        max_papers = DEFAULT_PAPERS
    max_papers = max(3, min(max_papers, MAX_PAPERS_CAP))

    style = (body.get("style") or "concise").strip().lower()
    if style not in ("concise", "comprehensive"):
        style = "concise"

    try:
        depth = int(body.get("depth", 1))
    except (TypeError, ValueError):
        depth = 1
    depth = max(1, min(depth, MAX_DEPTH))

    model = (body.get("model") or "").strip()
    if model not in ALLOWED_MODELS:
        model = ""  # fall back to the server default

    try:
        year_min = int(body.get("year_min", 0) or 0)
    except (TypeError, ValueError):
        year_min = 0
    if year_min and not (1990 <= year_min <= 2100):
        year_min = 0

    # Rate limit only counts when we actually start a new run (cache hits are free).
    cached = jobs._store.find_cached(
        topic, max_papers, depth, model, year_min, style, CACHE_TTL_SECONDS
    )
    if not cached:
        allowed, retry_in = limiter.allow(_client_ip(request))
        if not allowed:
            return JSONResponse(
                {"error": f"Rate limit reached. Try again in ~{retry_in // 60 + 1} min."},
                status_code=429,
            )

    job = jobs.submit(topic, max_papers, depth, model, year_min, style)
    return JSONResponse(
        {"job_id": job.id, "max_papers": max_papers, "depth": depth,
         "style": style, "cached": job.cached}
    )


@app.get("/api/review/{job_id}")
def get_review(job_id: str) -> JSONResponse:
    job = jobs.get(job_id)
    if not job:
        return JSONResponse({"error": "Unknown or expired job."}, status_code=404)
    return JSONResponse(
        {
            "id": job.id,
            "topic": job.topic,
            "status": job.status,
            "log": job.log,
            "review": job.review,
            "gaps": job.gaps,
            "papers": job.papers,
            "warnings": job.warnings,
            "model": job.model,
            "style": job.style,
            "cached": job.cached,
            "error": job.error,
            "stats": job.stats,
        }
    )


@app.get("/api/config")
def get_config() -> dict:
    return {
        "models": ALLOWED_MODELS,
        "max_papers": MAX_PAPERS_CAP,
        "default_papers": DEFAULT_PAPERS,
        "max_depth": MAX_DEPTH,
        "styles": ["concise", "comprehensive"],
    }


@app.get("/r/{job_id}", response_class=HTMLResponse)
def shared(job_id: str) -> HTMLResponse:
    # Same SPA; the page reads the job id from the path and loads it read-only.
    return index()


_static_dir = os.path.join(HERE, "static")
if os.path.isdir(_static_dir):
    app.mount("/static", StaticFiles(directory=_static_dir), name="static")