app.py

"""FastAPI web server for the DeepResearch pipeline.

Start with:
    uv run uvicorn app:app --reload --port 8000

Then open http://localhost:8000 in your browser.
"""
from __future__ import annotations

import asyncio
import json
import queue
import threading
from pathlib import Path
from typing import Any

from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, HTMLResponse
from pydantic import BaseModel

from src.storage import RunStore, RUNS_DIR
from src.schemas import ApprovedURLs, SearchResults, RunMeta, Report
from src.pipeline import run_pipeline, retry_run

app = FastAPI(title="DeepResearch")

# Per-run in-memory state (lost on server restart — expected for a dev tool)
_run_queues: dict[str, queue.Queue] = {}
_approval_events: dict[str, threading.Event] = {}
_approval_urls: dict[str, list[str]] = {}
_pending_candidates: dict[str, list[dict]] = {}


# ---------- Request/response models ----------

class StartRunRequest(BaseModel):
    query: str
    auto_approve: bool = False
    auto_approve_top: int = 20


class ApproveRequest(BaseModel):
    urls: list[str]


# ---------- Routes ----------

@app.get("/", response_class=HTMLResponse)
async def index() -> str:
    html = Path("static/index.html")
    if not html.exists():
        raise HTTPException(500, "static/index.html not found. Run from the project root.")
    return html.read_text(encoding="utf-8")


@app.get("/api/runs")
async def list_runs() -> list[dict]:
    store = RunStore()
    runs = store.list_runs(limit=50)
    return [m.model_dump(mode="json") for m in runs]


@app.get("/api/runs/{run_id}")
async def get_run(run_id: str) -> dict:
    store = RunStore()
    try:
        meta = store.load(run_id, "meta", RunMeta)
    except FileNotFoundError:
        raise HTTPException(404, f"Run {run_id!r} not found")

    # Recover zombie "running" runs (server restarted mid-run but report was saved)
    if meta.status == "running" and (RUNS_DIR / run_id / "report.md").exists():
        meta.status = "completed"

    result = meta.model_dump(mode="json")
    result["has_active_stream"] = run_id in _run_queues
    result["has_pending_approval"] = run_id in _pending_candidates
    return result


@app.get("/api/runs/{run_id}/report")
async def get_report(run_id: str) -> dict:
    report_path = RUNS_DIR / run_id / "report.md"
    if not report_path.exists():
        raise HTTPException(404, "Report not yet available")

    store = RunStore()
    try:
        meta = store.load(run_id, "meta", RunMeta)
    except FileNotFoundError:
        raise HTTPException(404, f"Run {run_id!r} not found")

    confidence = "unknown"
    try:
        report = store.load(run_id, "report", Report)
        confidence = report.confidence
    except Exception:
        pass

    return {
        "run_id": run_id,
        "markdown": report_path.read_text(encoding="utf-8"),
        "query": meta.query,
        "confidence": confidence,
        "sources_attempted": meta.sources_attempted,
        "sources_succeeded": meta.sources_succeeded,
        "generated_at": meta.completed_at.isoformat() if meta.completed_at else None,
    }


@app.get("/api/runs/{run_id}/candidates")
async def get_candidates(run_id: str) -> dict:
    """Return pending approval candidates (for reconnect after SSE drop)."""
    if run_id not in _pending_candidates:
        raise HTTPException(404, "No pending approval for this run")
    return {"candidates": _pending_candidates[run_id]}


@app.post("/api/runs")
async def start_run(req: StartRunRequest) -> dict:
    store = RunStore()
    run_id = store.create_run(req.query)

    q: queue.Queue = queue.Queue()
    approval_event = threading.Event()
    _run_queues[run_id] = q
    _approval_events[run_id] = approval_event

    def approval_fn(results: SearchResults) -> ApprovedURLs:
        candidates = [c.model_dump(mode="json") for c in results.candidates]
        _pending_candidates[run_id] = candidates
        q.put({"type": "approval_needed", "candidates": candidates})
        approval_event.wait()
        _pending_candidates.pop(run_id, None)
        return ApprovedURLs(urls=_approval_urls.get(run_id, []))

    def progress_fn(event: dict[str, Any]) -> None:
        if run_id in _run_queues:
            _run_queues[run_id].put(event)

    def thread_fn() -> None:
        try:
            run_pipeline(
                req.query,
                run_id=run_id,
                auto_approve=req.auto_approve,
                auto_approve_top=req.auto_approve_top,
                approval_fn=None if req.auto_approve else approval_fn,
                progress_fn=progress_fn,
            )
        except Exception as e:
            if run_id in _run_queues:
                try:
                    _run_queues[run_id].put({"type": "error", "message": str(e)})
                except Exception:
                    pass

    thread = threading.Thread(target=thread_fn, daemon=True)
    thread.start()

    return {"run_id": run_id}


@app.post("/api/runs/{run_id}/retry")
async def retry_run_endpoint(run_id: str) -> dict:
    """Resume a failed run from its last completed checkpoint.

    Re-uses the existing run_id and all cached artifacts. Sets up a fresh SSE
    queue and approval event so the UI can follow progress normally.
    """
    store = RunStore()
    try:
        meta = store.load(run_id, "meta", RunMeta)
    except FileNotFoundError:
        raise HTTPException(404, f"Run {run_id} not found")

    if meta.status == "running":
        raise HTTPException(400, "Run is already in progress")

    q: queue.Queue = queue.Queue()
    approval_event = threading.Event()
    _run_queues[run_id] = q
    _approval_events[run_id] = approval_event

    def approval_fn(results: SearchResults) -> ApprovedURLs:
        candidates = [c.model_dump(mode="json") for c in results.candidates]
        _pending_candidates[run_id] = candidates
        q.put({"type": "approval_needed", "candidates": candidates})
        approval_event.wait()
        _pending_candidates.pop(run_id, None)
        return ApprovedURLs(urls=_approval_urls.get(run_id, []))

    def progress_fn(event: dict[str, Any]) -> None:
        if run_id in _run_queues:
            _run_queues[run_id].put(event)

    def thread_fn() -> None:
        try:
            retry_run(run_id, approval_fn=approval_fn, progress_fn=progress_fn)
        except Exception as e:
            if run_id in _run_queues:
                try:
                    _run_queues[run_id].put({"type": "error", "message": str(e)})
                except Exception:
                    pass

    thread = threading.Thread(target=thread_fn, daemon=True)
    thread.start()

    return {"run_id": run_id}


@app.post("/api/runs/{run_id}/approve")
async def approve_run(run_id: str, req: ApproveRequest) -> dict:
    if run_id not in _approval_events:
        raise HTTPException(400, "No pending approval for this run (or server restarted)")
    if _approval_events[run_id].is_set():
        raise HTTPException(400, "Approval already submitted")

    _approval_urls[run_id] = req.urls
    _approval_events[run_id].set()
    return {"ok": True, "approved_count": len(req.urls)}


@app.get("/api/runs/{run_id}/events")
async def run_events(run_id: str) -> StreamingResponse:
    if run_id not in _run_queues:
        raise HTTPException(404, f"No active stream for run {run_id!r}. The run may have finished or the server restarted.")

    q = _run_queues[run_id]

    def _blocking_get() -> dict | None:
        try:
            return q.get(timeout=1.0)
        except queue.Empty:
            return None

    async def generate():
        loop = asyncio.get_event_loop()
        while True:
            event = await loop.run_in_executor(None, _blocking_get)
            if event is None:
                yield ": keepalive\n\n"
                continue
            yield f"data: {json.dumps(event)}\n\n"
            if event.get("type") in ("done", "error"):
                _run_queues.pop(run_id, None)
                _approval_events.pop(run_id, None)
                _approval_urls.pop(run_id, None)
                break

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
            "Connection": "keep-alive",
        },
    )