ui/chat_engine.py

"""
ui/chat_engine.py
-----------------
Chat-prep: turning a paper record into a retriever + chain, in three phases, with
SONIC's progress bar and pep-quotes animating over the slow bits.

`prepare_chat_stream()` is the public entrypoint. It's a generator so the Gradio
handler can relay each frame straight to the page.

The phase split exists because of ZeroGPU. Each phase wants something different:

  1. fetch      — network-bound. No GPU. Threaded, so the bar animates.
  2. vectorize  — the expensive encode. THE one GPU step (see ui/gpu.py).
  3. assemble   — re-open the persisted store on CPU + build the chain. Fast.

Phase 2 is the awkward one: ZeroGPU dispatches to its GPU worker from the calling
thread, so it must NOT run on a thread we spawned ourselves — which is precisely
what the original single-threaded-worker design did. On ZeroGPU we therefore call
it inline and the bar holds still for the few seconds it takes. Off ZeroGPU
there's no such constraint and the encode is slow (tens of seconds on CPU), so we
keep the old threaded animation. Same UX where it matters, correctness where it
counts.
"""

import random
import threading
import time

from ui.constants import SONIC_QUOTES
from ui.gpu import ON_ZEROGPU, vectorize_on_gpu
from ui.papers import resolve_and_download_pdf
from ui.sonic import SONIC_DATA_URI


def vec_loading_html(quote: str, phase: str, pct: int) -> str:
    return (
        f'<div class="vec-wrap">'
        f'  <div class="vec-figure"><img src="{SONIC_DATA_URI}" alt="SONIC"/></div>'
        f'  <div class="vec-quote"><span class="q">SONIC:</span> &ldquo;{quote}&rdquo;</div>'
        f'  <div class="vec-bar"><div class="vec-fill" style="width:{pct}%"></div></div>'
        f'  <div class="vec-phase">{phase}</div>'
        f'</div>'
    )


def _quote_at(quotes, start):
    """Rotate every ~4s, as the Streamlit build did."""
    return quotes[int((time.time() - start) // 4) % len(quotes)]


# --- phase workers (touch no UI) -------------------------------------------

def _download_blocking(record: dict, holder: dict):
    try:
        pdf_path = resolve_and_download_pdf(record)
        if not pdf_path:
            holder["error"] = ("Couldn't find a readable open-access PDF for this paper — it may be "
                               "paywalled. Use the 📄 Paper button to read it on the source site.")
            return
        holder["pdf_path"] = pdf_path
    except Exception as e:
        holder["error"] = f"Couldn't fetch this paper: {type(e).__name__}: {e}"


def _vectorize_blocking(pdf_path: str, holder: dict):
    try:
        vectorize_on_gpu(pdf_path)
    except Exception as e:
        holder["error"] = f"Couldn't read this paper: {type(e).__name__}: {e}"


def _assemble_session(pdf_path: str) -> dict:
    """Re-open the vectorstore on CPU and build the retriever + chain.

    Cheap: phase 2 already persisted the vectors, so build_vectorstore
    short-circuits to a plain load. Everything here is CPU-resident by design —
    it outlives the GPU window (see ui/gpu.py).
    """
    from vectorizeer import build_vectorstore
    from Qa import build_chain, get_llm, get_retriever
    vs = build_vectorstore(pdf_path)
    return {"retriever": get_retriever(vs), "chain": build_chain(get_llm())}


# --- the stream -------------------------------------------------------------

def prepare_chat_stream(record: dict):
    """Yield (html, done, session, error) frames while the paper is fetched and
    vectorized. Every frame before the last has done=False; the final frame
    carries either a session or an error."""
    quotes = random.sample(SONIC_QUOTES, len(SONIC_QUOTES))
    start = time.time()
    holder: dict = {}

    def fail(msg):
        return vec_loading_html(_quote_at(quotes, start), "…", 100), True, None, msg

    # 1. FETCH — threaded so the bar moves while the network does its thing.
    worker = threading.Thread(target=_download_blocking, args=(record, holder), daemon=True)
    worker.start()
    pct = 6
    # Emit one frame up front: a cached PDF can finish before the first is_alive()
    # check, and without this the panel would sit blank until the next phase.
    yield vec_loading_html(_quote_at(quotes, start), "Fetching the PDF…", pct), False, None, None
    while worker.is_alive():
        pct = min(pct + 2, 44)
        yield vec_loading_html(_quote_at(quotes, start), "Fetching the PDF…", pct), False, None, None
        time.sleep(0.35)
    worker.join()
    if holder.get("error"):
        yield fail(holder["error"])
        return
    pdf_path = holder["pdf_path"]

    # 2. VECTORIZE — the GPU step.
    if ON_ZEROGPU:
        # Must run on this thread: ZeroGPU hands the GPU to the caller, and a
        # thread we spawned isn't one. It's only a few seconds on a GPU, so the
        # bar simply holds rather than animating.
        yield vec_loading_html(_quote_at(quotes, start), "Reading & vectorizing every page…", 55), False, None, None
        _vectorize_blocking(pdf_path, holder)
        if holder.get("error"):
            yield fail(holder["error"])
            return
    else:
        # No such constraint on CPU — and here the encode is genuinely slow, so
        # the animation earns its keep.
        worker = threading.Thread(target=_vectorize_blocking, args=(pdf_path, holder), daemon=True)
        worker.start()
        while worker.is_alive():
            pct = min(pct + 2, 88)
            yield vec_loading_html(_quote_at(quotes, start),
                                   "Reading & vectorizing every page…", pct), False, None, None
            time.sleep(0.35)
        worker.join()
        if holder.get("error"):
            yield fail(holder["error"])
            return

    # 3. ASSEMBLE — CPU, fast.
    yield vec_loading_html(_quote_at(quotes, start), "Almost there…", 94), False, None, None
    try:
        session = _assemble_session(pdf_path)
    except Exception as e:
        yield fail(f"Couldn't prepare this paper for chat: {type(e).__name__}: {e}")
        return

    yield vec_loading_html(quotes[0], "Ready.", 100), True, session, None