app/main.py

"""TorchDocs Agent — Gradio web app (M5).

A long-lived server: the embedding model loads once at startup, so each
question is answered in seconds (unlike the batch Actions runs). Ask a PyTorch
question in English; the agent searches the docs and answers with clickable
citations.

Concurrent by default: each question is answered from request-local state
(agent/loop.py builds fresh sections/transcript/budgets per call), so many
users can be served at once. The Gradio queue is opened to TORCHDOCS_CONCURRENCY
workers instead of the framework's serial default — the work is almost all I/O
(LLM + Neon), so it overlaps cleanly and nobody waits in line.

Run locally:  python -m app.main   (needs NEON_URL + OpenRouter env, see .env)
Deploy:       Hugging Face Spaces (this file is the Space entrypoint).
"""

from __future__ import annotations

import html
import itertools
import os
import threading
import time
from collections import defaultdict, deque

import gradio as gr
from dotenv import load_dotenv

from agent.route import answer_routed
from agent.schemas import Answer

load_dotenv()

INTRO = (
    "# 🔥 TorchDocs Agent\n"
    "Ask anything about PyTorch — in English. Answers are grounded in the "
    "official documentation with clickable citations; source-code questions are "
    "referred to GitHub / DeepWiki."
)

EXAMPLES = [
    "How do I use torch.optim.SGD with momentum?",
    "What LR schedulers are supported in PyTorch?",
    "How do I build a CNN to classify images, end to end?",
    "How is conv2d implemented under the hood?",
]

# Shown under the citations: a link to the PyTorch license, its name as the text.
LICENSE_NOTE = "<sub>[BSD-3-Clause](https://github.com/pytorch/pytorch/blob/main/LICENSE)</sub>"

# How many questions to answer at once. The default is generous because a
# request spends nearly all its wall-clock waiting on the LLM and Neon, not on
# CPU — overlapping them is what turns "wait your turn" into "answered now".
# Override per deploy (a bigger Space, a paid LLM key) via the env var.
CONCURRENCY = int(os.environ.get("TORCHDOCS_CONCURRENCY", "16"))

# Backpressure: how many requests may WAIT behind the concurrent workers before
# new ones are turned away. Without a cap the queue grows without bound under a
# flood, and everyone in it waits forever instead of being told "busy".
QUEUE_SIZE = int(os.environ.get("TORCHDOCS_QUEUE_SIZE", "64"))

# Per-client throttle: at most RATE_LIMIT questions per RATE_WINDOW seconds per
# client IP, so one over-eager caller can't occupy every worker slot (and burn
# the shared free-tier LLM quota) by itself. 0 disables the throttle.
RATE_LIMIT = int(os.environ.get("TORCHDOCS_RATE_LIMIT", "8"))
RATE_WINDOW = float(os.environ.get("TORCHDOCS_RATE_WINDOW_SECONDS", "60"))
BUSY_NOTE = "You're asking faster than I can answer — give it a moment and try again."

# Shown the instant a question is submitted, then replaced by the answer. The
# heavy path (guard embed → retrieval → LLM) takes a few seconds, so immediate
# feedback is the difference between "is it broken?" and "it's working".
THINKING_NOTE = "🔎 Searching the PyTorch docs…"
# We can't cheaply stream answer TOKENS — the answer is a validated JSON object
# assembled over several tool calls, not free prose. So instead we stream the
# REASONING: the pipeline emits a short trace line per step (which docs it
# searched, what it found, when it starts writing) and respond() renders those
# in grey with a turning wheel, then replaces them with the answer in normal
# (black) text. A multi-second wait then reads as visible work, not a freeze.
THINKING_SPINNER = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏"
THINKING_TICK = 0.6  # seconds between animation frames
# grey trace lines take the theme's subdued colour (adapts to light/dark); the
# inline style survives Gradio's markdown sanitiser (verified on gradio 6.20)
TRACE_STYLE = "color:var(--body-text-color-subdued)"
# Keep the phrase "went wrong" — the post-deploy smoke test treats it as the
# failure marker (scripts/smoke_space.py). The real exception goes to the logs;
# the user never sees hosts, model slugs, or config internals.
ERROR_NOTE = "⚠️ Something went wrong answering that. Please try again in a moment."

_RATE_LOCK = threading.Lock()
_RATE_BUCKETS: dict[str, deque[float]] = defaultdict(deque)


def _rate_limited(client_id: str) -> bool:
    """Sliding window: True if this client already used its RATE_LIMIT slots."""
    now = time.monotonic()
    with _RATE_LOCK:
        bucket = _RATE_BUCKETS[client_id]
        while bucket and now - bucket[0] > RATE_WINDOW:
            bucket.popleft()
        if len(bucket) >= RATE_LIMIT:
            return True
        bucket.append(now)
        if len(_RATE_BUCKETS) > 4096:  # keep one-off visitors from growing the table
            for key in [k for k, b in _RATE_BUCKETS.items() if not b or now - b[-1] > RATE_WINDOW]:
                del _RATE_BUCKETS[key]
        return False


def _warm_up() -> None:
    """Load the embedding model once so the first question isn't slow.

    This also covers the guard: its topicality check embeds the question with
    the same model.
    """
    try:
        from index.embed import embed_query

        embed_query("warmup")
    except Exception as exc:  # noqa: BLE001 — warmup is best-effort
        print(f"[app] warmup skipped: {exc}")


def render(answer: Answer) -> str:
    """Answer + citations + referrals as one markdown block."""
    parts = [answer.answer_md]
    if answer.citations:
        parts.append("\n---\n**Sources**")
        for c in answer.citations:
            frag = f"#{c.anchor}" if c.anchor else ""
            label = c.title or c.url
            parts.append(f"- [{label}]({c.url}{frag})")
    if answer.referrals:
        parts.append("\n**Beyond these docs**")
        for r in answer.referrals:
            parts.append(f"- [{r.reason or r.url}]({r.url})")
    if answer.torch_version and answer.torch_version != "unknown":
        parts.append(f"\n<sub>targets PyTorch {answer.torch_version}</sub>")
    if answer.citations:  # only when we actually quoted documentation
        parts.append("\n" + LICENSE_NOTE)
    return "\n".join(parts)


def _pipeline(
    question: str,
    request: gr.Request = None,
    out: dict | None = None,
    progress=None,
) -> str:
    """The full answer pipeline → final markdown string (no UI concerns).

    `out` (optional) receives the Answer object under "answer" when one was
    generated — respond()'s freshness pass needs the citations, and returning a
    tuple would break every caller that treats the result as markdown.

    `progress` (optional) is a sink for short trace lines the pipeline emits as
    it retrieves and reasons; respond() streams them to the UI in grey.
    """
    question = (question or "").strip()
    if not question:
        return "Ask me something about PyTorch."
    # gradio injects `request` for real traffic; direct calls (tests) skip it
    client = getattr(getattr(request, "client", None), "host", None)
    if client and RATE_LIMIT > 0 and _rate_limited(client):
        return BUSY_NOTE
    from agent.guard import guard

    verdict = guard(question)  # one check on the raw user input, before the pipeline
    if not verdict.ok:
        return verdict.message
    try:
        # routed: simple questions take the 1-2-call grounded path (seconds),
        # multi-source shapes get the full tool loop (see agent/route.py)
        started = time.monotonic()
        answer = answer_routed(question, progress=progress)
        # question→answer latency is the core UX metric — log it per request so
        # the Space logs show real p50/p95, not just the eval's sampled number
        print(f"[app] answered in {time.monotonic() - started:.1f}s", flush=True)
        if out is not None:
            out["answer"] = answer
        return render(answer)
    except Exception as exc:  # noqa: BLE001 — never crash the UI
        # the real error goes to the logs; the user gets a generic line, since
        # an exception string can leak hosts, model slugs, and config internals
        print(f"[app] answer failed: {type(exc).__name__}: {exc}", flush=True)
        return ERROR_NOTE


def _render_trace(lines: list[str], spinner: str | None) -> str:
    """The live reasoning trace as one grey markdown block.

    Each step on its own line, escaped (a step echoes the user's query, which is
    untrusted). While work continues, `spinner` is a turning wheel on a trailing
    line; pass None to drop it. Empty trace + spinner → just the wheel.
    """
    rows = [html.escape(line) for line in lines]
    if spinner is not None:
        rows.append(spinner)
    if not rows:
        rows = ["…"]
    return f'<span style="{TRACE_STYLE}">{"<br>".join(rows)}</span>'


# The most recent background freshness thread, kept only as a test seam so a
# test can join it deterministically instead of sleeping. Production ignores it.
_LAST_FRESHNESS: threading.Thread | None = None


def _spawn_freshness(urls: list[str]) -> threading.Thread | None:
    """Start the post-answer freshness heal on a detached daemon thread.

    Returns the started thread (or None when freshness is disabled). The UI has
    already delivered its final answer, so this only ever improves the NEXT
    answer — it must not block or stream. Best-effort: every failure is caught
    and logged on the worker so it can never surface to a user.
    """
    from index import freshness

    if not freshness.enabled():
        return None

    def _run() -> None:
        try:
            freshness.refresh_pages(urls)
        except Exception as exc:  # noqa: BLE001 — background heal is best-effort
            print(f"[app] background freshness failed: {type(exc).__name__}: {exc}", flush=True)

    thread = threading.Thread(target=_run, daemon=True)
    thread.start()
    global _LAST_FRESHNESS
    _LAST_FRESHNESS = thread
    return thread


def respond(question: str, request: gr.Request = None):
    """UI entrypoint: show a LIVE thinking indicator, then the answer.

    A generator so Gradio streams feedback the instant a question is submitted.
    The pipeline runs on a worker thread while this generator emits an animated
    spinner + stage label every THINKING_TICK seconds — so the multi-second wait
    reads as "working", not "frozen" — then yields the finished markdown as its
    LAST value and ends. The spinner ends WITH the answer: there is no
    post-answer spinner.

    Freshness self-heal (index/freshness.py) still runs, but entirely in the
    BACKGROUND — a detached daemon thread revalidates the cited pages against
    the live docs and heals any drifted chunks in place, so the NEXT asker of
    this question gets the corrected answer. It never touches this stream:
    revalidation is a live, multi-page network fetch, and holding the user's
    spinner on it spun the wheel for minutes when the docs server was slow.
    Best-effort — any failure is swallowed on the worker thread.

    The first yield is the static THINKING_NOTE (immediate paint);
    gradio_client.predict returns the LAST yielded value, so the smoke test
    still gets a real answer.
    """
    yield THINKING_NOTE

    result: dict = {}
    trace: list[str] = []
    partial: dict = {"md": ""}  # the answer prose as it streams in
    trace_lock = threading.Lock()  # the worker appends; this generator reads

    def progress(line: str) -> None:
        with trace_lock:
            trace.append(line)

    def on_delta(prose: str) -> None:
        with trace_lock:
            partial["md"] = prose

    def work():
        # set the streaming sink INSIDE the worker thread: the answer pipeline
        # runs here, and a new thread has its own context, so the compat answer
        # path (agent/llm.py) sees the sink and streams answer_md tokens through
        # on_delta. Reset on the way out so nothing leaks across requests.
        from agent import llm

        token = llm.answer_stream_sink.set(on_delta)
        try:
            result["md"] = _pipeline(question, request, out=result, progress=progress)
        except Exception as exc:  # noqa: BLE001 — the UI must never hang on a crash
            print(f"[app] pipeline thread failed: {type(exc).__name__}: {exc}", flush=True)
            result["md"] = ERROR_NOTE
        finally:
            llm.answer_stream_sink.reset(token)

    worker = threading.Thread(target=work, daemon=True)
    worker.start()
    frames = itertools.cycle(THINKING_SPINNER)
    while worker.is_alive():
        worker.join(timeout=THINKING_TICK)  # wait a tick (or finish sooner)
        with trace_lock:
            lines = list(trace)
            prose = partial["md"]
        # once the answer starts streaming, show the prose growing (with a
        # block cursor) instead of the grey trace — the wait turns into the
        # answer being written. Before that, the animated trace + wheel.
        if prose:
            yield f"{prose} ▍"
        else:
            yield _render_trace(lines, next(frames))  # grey trace + turning wheel
    md = result.get("md", ERROR_NOTE)
    yield md  # the finished answer, in normal (black) text — the generator ends
    # here. No post-answer spinner: the wheel stopped the instant the answer
    # appeared. The freshness self-heal below is detached and never streams.

    # Fire-and-forget freshness: revalidate the CITED pages against the live
    # docs and heal any drifted chunks in place, so the NEXT asker of this
    # question gets the corrected answer. Background-only by decision — a live,
    # multi-page network fetch must never hold the user's spinner (it used to,
    # and a slow docs server spun it for minutes). Best-effort; the worker
    # swallows and logs any failure, and the shown answer is already final.
    answer = result.get("answer")
    if answer is not None and answer.citations:
        _spawn_freshness([c.url for c in answer.citations])


def build_ui():
    with gr.Blocks(title="TorchDocs Agent") as demo:
        gr.Markdown(INTRO)
        # lines=1 (NOT 2) is load-bearing: Gradio only submits on a bare Enter for
        # a SINGLE-line textbox — a multi-line box (lines>1) treats Enter as a
        # newline and submits on Shift+Enter instead. max_lines still lets a long
        # question grow visually; the submit rule keys off the `lines` prop, not
        # the rendered height, so Enter keeps sending. Don't bump lines back to 2.
        question = gr.Textbox(
            label="Your question",
            placeholder="How do I use a DataLoader?",
            lines=1,
            max_lines=6,
        )
        ask = gr.Button("Ask", variant="primary")
        answer = gr.Markdown(label="Answer")
        gr.Examples(EXAMPLES, inputs=question)
        # api_name gives the post-deploy smoke test (scripts/smoke_space.py) a
        # stable Gradio endpoint to call: client.predict(..., api_name="/respond")
        ask.click(respond, inputs=question, outputs=answer, api_name="respond")
        question.submit(respond, inputs=question, outputs=answer, api_name=False)
    return demo


def serve(demo) -> None:
    """Launch the UI with the deployment bind settings (shared by both entrypoints).

    Opening the queue is what makes the app concurrent: Gradio 4/5 default every
    event to a serial `concurrency_limit=1`, so without this each user waits for
    the previous answer to finish. `max_threads` is lifted in step so the worker
    pool — which also holds threads parked in retry back-off on a 429 — never
    becomes the hidden ceiling below CONCURRENCY. `max_size` bounds how many
    requests may wait behind them, so a flood gets "queue full" instead of an
    ever-growing line.
    """
    demo.queue(default_concurrency_limit=CONCURRENCY, max_size=QUEUE_SIZE)
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        max_threads=max(40, CONCURRENCY * 2),
    )


def main() -> None:
    _warm_up()
    serve(build_ui())


if __name__ == "__main__":
    main()