app.py

"""Gradio Server Mode backend for the Repair Guy Space.

A hands-busy mechanic's assistant: a page viewer driven by short requests. It
finds and points — it never writes answers. All ingestion happens offline via
scripts/index_modal.py; the Space only syncs the pre-indexed library and serves
turns.

  Obvious nav (client)  next/previous page, "page 412", back, next/previous
                        section: pure frontend state, no server call.
  Everything else       → /find, one ZeroGPU call (pipelines/agent_ask.py):
  (ZeroGPU agent turn)  MiniCPM5-1B (the text "brain") sees the conversation so
                        far and the WHOLE text of the page being viewed, and
                        calls tools in a loop until a page is shown — circle on
                        the current page (MiniCPM-V grounds the box), jump to a
                        page, or search. No table of contents is injected; the
                        agent navigates by retrieval and the current page's text.
                        Search is FUSED: ColEmbed (visual store) shortlists pages,
                        the 1B reranks by their parsed text. Streamed as events;
                        the UI shows tool chips and the resulting page/circle,
                        never model prose. History is kept only to resolve
                        references.

The frontend's breadcrumb / next-previous-section navigation uses the manual's
clean PDF-bookmark chapters (served by /sections/{doc_id}); the agent turn does
not receive them.

UI architecture — this is NOT a gr.Blocks app. It runs in Gradio *Server Mode*
(`gradio.Server`, a FastAPI server with Gradio's engine: queueing, streaming
and — crucially — ZeroGPU auth). The Python pipelines below are exposed as
`@app.api` endpoints; the frontend is a fully custom single page (Tailwind +
Alpine, no build step) served from frontend/index.html, which calls those
endpoints with @gradio/client so the HF iframe auth headers ZeroGPU needs are
forwarded. The source PDFs are served straight from FastAPI at /pdf/<doc_id>.

Module layout:
  models/colembed.py        ColEmbed       — search shortlist: page embeddings + MaxSim
  models/minicpm_agent.py   MiniCPM5-1B    — the agent "brain": tool loop + rerank
  models/minicpm.py         MiniCPM-V      — the "eyes": grounds the circle
  models/nemotron_embed.py  NemotronEmbed  — parsed: dense chunk/query embeddings (ingest)
  core/visual_store.py      VisualStore    — on-disk per-page token embeddings
  core/parsed_store.py      ParsedStore    — chunks, embeddings, raw parsed pages
  core/page_context.py      whole-page → text for the agent's context
  pipelines/agent_ask.py    AgentPipeline — the agent find-and-point GPU turn
  pipelines/mock_ask.py     MockAskPipeline — local UI iteration (MOCK_MODELS=1)
  frontend/index.html       the custom UI
"""

import base64
import io
import json
import logging
import os
import shutil
import time
import warnings

import gradio as gr
from fastapi.responses import FileResponse, HTMLResponse, Response
from huggingface_hub import HfApi, snapshot_download

from core.constants import (
    AGENT_MODELS,
    DEFAULT_AGENT_MODEL,
    DEFAULT_RETRIEVAL_MODE,
    DEFAULT_TOP_K,
    GROUND_ENABLE_THINKING,
    LIBRARY_DATASET_ID,
    MAX_TOP_K,
    MOCK_MODELS,
    MOCK_PDF_DIR,
    PARSED_SUBDIR,
    PREINDEXED_DIR,
    RENDER_DPI,
    RETRIEVAL_MODES,
    VISUAL_SUBDIR,
)
from core.pdf import pdf_outline, render_page_png

# gradio 6.17.3 (pinned — see README frontmatter) still uses starlette's old
# 422 constant, so every queue join emits a StarletteDeprecationWarning. Not
# ours to fix; silence it so the Space logs stay readable.
warnings.filterwarnings(
    "ignore", message=r"'HTTP_422_UNPROCESSABLE_ENTITY' is deprecated"
)

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s | %(message)s"
)
logging.getLogger("httpx").setLevel(logging.WARNING)  # logs every hub request
log = logging.getLogger("repairguy")


def _build_libraries():
    """(visual_store, parsed_store, agent_pipeline), constructed once at startup.

    The agent path is FUSED: ColEmbed (visual store) supplies the search
    shortlist, the parsed store supplies the page text the 1B reranks with and
    the agent reasons over — so both stores are live at once.

    In MOCK_MODELS mode a single PDF-folder mock backs both stores and no
    GPU/model code is imported (those modules load CUDA at import). Otherwise the
    real stores and the agent pipeline load — the models onto cuda here, in the
    main process."""
    if MOCK_MODELS:
        from pipelines.mock_ask import MockAskPipeline, MockStore

        store = MockStore()
        print(f"⚠️  MOCK_MODELS — fake agent over PDFs in {MOCK_PDF_DIR}")
        return store, store, MockAskPipeline()

    from core.parsed_store import ParsedStore
    from core.visual_store import VisualStore
    from pipelines.agent_ask import AgentPipeline

    return (
        VisualStore(os.path.join(PREINDEXED_DIR, VISUAL_SUBDIR)),
        ParsedStore(os.path.join(PREINDEXED_DIR, PARSED_SUBDIR)),
        AgentPipeline(),
    )


VISUAL_STORE, PARSED_STORE, PIPELINE = _build_libraries()
# Valid agent-brain keys, for validating the per-request `agent_model`.
_AGENT_MODEL_KEYS = {m["key"] for m in AGENT_MODELS}
# Brains too large to keep the ColEmbed visual retriever resident alongside them
# (constants.AGENT_MODELS `forbid_visual`). While one is active, a "visual"
# retrieval_mode is forced to "parsed" so ColEmbed never loads next to it (OOM).
_NO_VISUAL_MODEL_KEYS = {m["key"] for m in AGENT_MODELS if m.get("forbid_visual")}
# Valid search-index keys, for validating the per-request `retrieval_mode`.
_RETRIEVAL_MODE_KEYS = {m["key"] for m in RETRIEVAL_MODES}
# method -> store, for the picker / pdf lookups. In mock both keys map to the
# one MockStore, so every mock manual reads as indexed under both methods.
_METHOD_STORES = {"visual": VISUAL_STORE, "parsed": PARSED_STORE}


def sync_library() -> None:
    """Pull pre-indexed manuals from the library dataset into /data.
    A missing or empty dataset just means an empty library."""
    if MOCK_MODELS:  # mock library is the local PDF folder; nothing to sync
        return
    try:
        snapshot_download(
            LIBRARY_DATASET_ID, repo_type="dataset", local_dir=PREINDEXED_DIR
        )
        remote_files = HfApi().list_repo_files(
            LIBRARY_DATASET_ID, repo_type="dataset"
        )
        log.info("library synced from %s", LIBRARY_DATASET_ID)
    except Exception as e:
        log.warning("library dataset not synced (%s): %s", LIBRARY_DATASET_ID, e)
        return
    # PREINDEXED_DIR mirrors the dataset (one dir per method); prune top-level
    # leftovers from the pre-method-prefix layout, which snapshot_download
    # never deletes.
    for entry in os.listdir(PREINDEXED_DIR):
        path = os.path.join(PREINDEXED_DIR, entry)
        if entry.startswith(".") or entry in (VISUAL_SUBDIR, PARSED_SUBDIR):
            continue
        if os.path.isdir(path) and os.path.isfile(os.path.join(path, "index.json")):
            print(f"Pruning stale pre-migration doc dir: {entry}")
            shutil.rmtree(path, ignore_errors=True)
    # snapshot_download only adds/updates; it never removes a local doc dir that
    # was deleted from the dataset. Prune any method/doc dir absent remotely so
    # deletions in the library propagate to the Space on the next sync.
    remote_docs = {
        (parts[0], parts[1])
        for f in remote_files
        if len(parts := f.split("/")) >= 3 and parts[0] in (VISUAL_SUBDIR, PARSED_SUBDIR)
    }
    for method in (VISUAL_SUBDIR, PARSED_SUBDIR):
        method_dir = os.path.join(PREINDEXED_DIR, method)
        if not os.path.isdir(method_dir):
            continue
        for doc_id in os.listdir(method_dir):
            doc_dir = os.path.join(method_dir, doc_id)
            if doc_id.startswith(".") or not os.path.isdir(doc_dir):
                continue
            if (method, doc_id) not in remote_docs:
                print(f"Pruning doc removed from library: {method}/{doc_id}")
                shutil.rmtree(doc_dir, ignore_errors=True)


sync_library()


def _manual_choices() -> list[dict]:
    """The manuals shown in the picker (doc ids are name slugs, so the same
    manual lands on the same id in both stores). The fused agent needs BOTH
    indexes; a manual missing one is tagged and the agent turn rejects it.
    pages drives the viewer's page count."""
    docs: dict[str, dict] = {}
    for method, store in _METHOD_STORES.items():
        for d in store.list_docs():
            entry = docs.setdefault(
                d["doc_id"], {"name": d["name"], "methods": set(), "pages": 0}
            )
            entry["methods"].add(method)
            entry["pages"] = max(entry["pages"], d["pages"])
    choices = []
    for doc_id, info in sorted(docs.items(), key=lambda kv: kv[1]["name"].lower()):
        label = info["name"]
        if info["methods"] != {"visual", "parsed"}:
            label += " — needs reindex"
        choices.append({"value": doc_id, "label": label, "pages": info["pages"]})
    return choices


def _pdf_path(doc_id: str) -> str | None:
    """The source PDF for a manual (kept identically in whichever store indexed
    it — both copy doc.pdf at ingest)."""
    for store in _METHOD_STORES.values():
        if store.exists(doc_id):
            return store.pdf_path(doc_id)
    return None


# The manual's own clean bookmark chapters (pdf_outline), cached per doc and
# cleared on library re-sync. Drives ONLY the frontend's breadcrumb +
# next/previous-section navigation (served by /sections/{doc_id}); the agent
# turn no longer receives a table of contents — it works from retrieval and the
# current page's text alone. A visual-only manual with no PDF bookmarks has no
# outline; section navigation then no-ops.
_OUTLINE_CACHE: dict[str, list[dict]] = {}


def _doc_outline(doc_id: str) -> list[dict]:
    if doc_id not in _OUTLINE_CACHE:
        if MOCK_MODELS:
            _OUTLINE_CACHE[doc_id] = PARSED_STORE.sections(doc_id)
        else:
            path = _pdf_path(doc_id)
            _OUTLINE_CACHE[doc_id] = pdf_outline(path) if path else []
    return _OUTLINE_CACHE[doc_id]


def _thumb_data_uri(img, width: int = 280) -> str:
    """A downscaled JPEG data URI for a rendered page, small enough to ship in
    the JSON answer (cited-page thumbnails are decorative, not full-res)."""
    img = img.convert("RGB")
    if img.width > width:
        img = img.resize((width, round(img.height * width / img.width)))
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=80)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()


# ---------------------------------------------------------------------------
# Server Mode: API engine (ZeroGPU auth, queueing) with a custom frontend.
# ---------------------------------------------------------------------------
app = gr.Server()


@app.api(name="manuals")
def api_manuals() -> list[dict]:
    """The manual picker's options: [{value, label}]."""
    return _manual_choices()


@app.api(name="refresh")
def api_refresh() -> list[dict]:
    """Re-sync the library dataset (incremental) and return the refreshed
    picker options, so manuals indexed after boot show up without a restart."""
    sync_library()
    _OUTLINE_CACHE.clear()
    return _manual_choices()


@app.api(name="find")
def api_find(
    request: str,
    manual: str = "",
    k: int = DEFAULT_TOP_K,
    page: int = 0,
    section: str = "",
    pages: list = None,
    history: list = None,
    think: bool = GROUND_ENABLE_THINKING,
    agent_model: str = DEFAULT_AGENT_MODEL,
    retrieval_mode: str = DEFAULT_RETRIEVAL_MODE,
    vram_log: bool = False,
    session_id: str = "",
) -> dict:  # the per-yield type: Server.api infers outputs from this annotation
    """One agent turn (one ZeroGPU call), streamed as events (see
    pipelines/agent_ask.py for the protocol). page/section are what the viewer
    currently shows (the agent's current page); history is the compact memory of
    past turns ([{request, action}]) for resolving references.

    Yields {type: status|step|tool_result|found|done|error, ...}; tool_result
    galleries are converted to JSON-able thumbnails here, and the terminal `done`
    carries elapsed/k. Soft errors so the frontend renders them as chips."""
    request = (request or "").strip()
    if not request:
        yield {"type": "error", "error": "Tell me what to find."}
        return
    if not manual:
        yield {"type": "error", "error": "Pick a manual first ☝️"}
        return
    if not (VISUAL_STORE.exists(manual) and PARSED_STORE.exists(manual)):
        yield {
            "type": "error",
            "error": "This manual needs reindexing — the assistant needs both the "
            "visual and parsed indexes. Pick another manual, or re-run indexing.",
        }
        return
    start = time.monotonic()
    # `pages` = every page currently on the viewer (the two-page spread shows
    # two); the agent reads them all and may circle on any. Falls back to the
    # single `page` for older clients.
    shown = [int(p) for p in (pages or []) if str(p).strip().isdigit()] or (
        [int(page)] if page else []
    )
    viewer = {"page": int(page or 0), "section": str(section or ""), "pages": shown}
    # Unknown model key → default (the pipeline falls back too; validate here so
    # the log reflects what actually ran).
    if agent_model not in _AGENT_MODEL_KEYS:
        agent_model = DEFAULT_AGENT_MODEL
    if retrieval_mode not in _RETRIEVAL_MODE_KEYS:
        retrieval_mode = DEFAULT_RETRIEVAL_MODE
    # The big brains can't share VRAM with the ColEmbed visual retriever — force
    # parsed so ColEmbed never loads alongside them (the UI also greys it out).
    if retrieval_mode == "visual" and agent_model in _NO_VISUAL_MODEL_KEYS:
        retrieval_mode = DEFAULT_RETRIEVAL_MODE
    log.info(
        "find: manual=%s k=%s think=%s model=%s search=%s viewer=%s hist=%d q=%r",
        manual, k, bool(think), agent_model, retrieval_mode, viewer,
        len(history or []), request[:200],
    )
    try:
        events = PIPELINE.run_find(
            VISUAL_STORE, PARSED_STORE, request, [manual], int(k),
            viewer, history, bool(think), agent_model, retrieval_mode,
            bool(vram_log), str(session_id or ""),
        )
        for ev in events:
            if ev.get("type") == "tool_result" and "gallery" in ev:
                pages = [p for _, p in ev["page_refs"]]
                yield {
                    "type": "tool_result",
                    "tool": ev["tool"],
                    "pages": pages,
                    "thumbnails": [
                        {"page": p, "src": _thumb_data_uri(img), "caption": cap}
                        for (img, cap), p in zip(ev["gallery"], pages)
                    ],
                }
            elif ev.get("type") == "done":
                elapsed = round(time.monotonic() - start, 1)
                log.info("find: done in %.1fs (%s)", elapsed, ev.get("kind"))
                yield {**ev, "elapsed": elapsed, "k": int(k)}
            else:
                yield ev
    except ValueError as e:
        log.warning("find: rejected — %s", e)
        yield {"type": "error", "error": f"⚠️ {e}"}
    except Exception as e:
        log.exception("find: failed after %.1fs", time.monotonic() - start)
        yield {"type": "error", "error": f"⚠️ Something went wrong: {e}"}


# --- custom FastAPI routes: serve the SPA and the source PDFs ---------------
_FRONTEND_DIR = os.path.join(os.path.dirname(__file__), "frontend")


@app.get("/")
def index():
    """Serve the single-page UI, injecting the small bit of server config the
    frontend needs (default/max k, default grounding-thinking, the agent-model
    list + default) so it needs no extra round-trip on load."""
    with open(os.path.join(_FRONTEND_DIR, "index.html")) as f:
        html = f.read()
    html = (
        html.replace("__DEFAULT_K__", str(DEFAULT_TOP_K))
        .replace("__MAX_K__", str(MAX_TOP_K))
        # Initial state of the settings-panel "thinking" toggle (a JS bool).
        .replace("__GROUND_THINK__", "true" if GROUND_ENABLE_THINKING else "false")
        # Agent-brain picker: the selectable models (key+label) and the default,
        # so the settings dropdown needs no extra round-trip on load.
        .replace(
            "__AGENT_MODELS_JSON__",
            json.dumps([
                {"key": m["key"], "label": m["label"],
                 "forbidVisual": bool(m.get("forbid_visual"))}
                for m in AGENT_MODELS
            ]),
        )
        .replace("__AGENT_MODEL__", DEFAULT_AGENT_MODEL)
        # Search-index picker: which index the search tool ranks against
        # (visual / parsed) and the default, so the settings dropdown needs no
        # extra round-trip on load.
        .replace(
            "__RETRIEVAL_MODES_JSON__",
            json.dumps([{"key": m["key"], "label": m["label"]} for m in RETRIEVAL_MODES]),
        )
        .replace("__RETRIEVAL_MODE__", DEFAULT_RETRIEVAL_MODE)
        # Cache-bust key for /page images: changing the render DPI changes the
        # served page size, so it must change the URL too — otherwise a browser
        # could keep an old-resolution page (cached up to a day) under the same
        # URL while the grounding render uses the new size.
        .replace("__PAGE_V__", str(RENDER_DPI))
    )
    # Never cache the SPA shell: it's tiny and inlines all the app JS, so a
    # cached copy silently runs stale code after a deploy (which masked an
    # earlier fix). Versioned assets (/page, /vendor) keep their long caches.
    return HTMLResponse(html, headers={"Cache-Control": "no-store"})


@app.get("/media/{filename}")
def serve_media(filename: str):
    """Static frontend assets (logo etc.) from frontend/assets/."""
    path = os.path.join(_FRONTEND_DIR, "assets", os.path.basename(filename))
    if not os.path.isfile(path):
        return Response(status_code=404)
    return FileResponse(path)


# Self-hosted browser assets for the voice mode: Silero VAD (the always-on
# listener) + onnxruntime-web's wasm runtime, vendored under frontend/vendor/ so
# nothing on the listening path depends on a CDN. Pinned + immutable → long
# cache. Only known extensions are served, and the resolved path is confined to
# the vendor dir (no traversal).
_VENDOR_DIR = os.path.join(_FRONTEND_DIR, "vendor")
_VENDOR_MEDIA_TYPES = {
    ".js": "text/javascript",
    ".mjs": "text/javascript",
    ".wasm": "application/wasm",
    ".onnx": "application/octet-stream",
    ".json": "application/json",  # Moonshine config/tokenizer (self-hosted ASR)
}


@app.get("/vendor/{path:path}")
def serve_vendor(path: str):
    """Vendored voice-mode assets (vad worklet/model + ort wasm runtime)."""
    full = os.path.normpath(os.path.join(_VENDOR_DIR, path))
    if not full.startswith(_VENDOR_DIR + os.sep) or not os.path.isfile(full):
        return Response(status_code=404)
    media_type = _VENDOR_MEDIA_TYPES.get(os.path.splitext(full)[1].lower())
    if media_type is None:
        return Response(status_code=404)
    return FileResponse(
        full,
        media_type=media_type,
        headers={"Cache-Control": "public, max-age=86400"},
    )


@app.get("/pdf/{doc_id}")
def serve_pdf(doc_id: str):
    """A manual's source PDF (kept for direct download/open-in-tab; the viewer
    pane shows /page images, not the PDF)."""
    path = _pdf_path(doc_id)
    if not path or not os.path.isfile(path):
        return Response(status_code=404)
    return FileResponse(
        path, media_type="application/pdf", content_disposition_type="inline"
    )


@app.get("/page/{doc_id}/{page}")
def serve_page(doc_id: str, page: int):
    """One manual page rendered to PNG at RENDER_DPI — the viewer pane's <img>
    source. Cached long (immutable content → instant page flips); the URL
    carries a ?v=RENDER_DPI cache key so a render-size change can't leave a
    stale-resolution page under the same URL. The circle overlay sizes its
    viewBox from the grounding render's dimensions (sent with the bbox), not
    this image, so it stays aligned even if a cached page is a different size."""
    path = _pdf_path(doc_id)
    if not path or not os.path.isfile(path):
        return Response(status_code=404)
    try:
        data = render_page_png(path, page)
    except ValueError:
        return Response(status_code=404)
    return Response(
        data,
        media_type="image/png",
        headers={"Cache-Control": "public, max-age=86400"},
    )


@app.get("/sections/{doc_id}")
def serve_sections(doc_id: str):
    """The manual's clean chapter outline [{title, page_start, page_end}] —
    drives the breadcrumb and next/previous-section navigation in the
    frontend. Empty for a visual-only manual with no PDF bookmarks."""
    return {"sections": _doc_outline(doc_id)}


if __name__ == "__main__":
    app.launch()