nova_app.py

"""
NOVA — Research, guided by SONIC
================================================================
A single Gradio app that stitches together the two projects, unchanged:

  • app/            the research pipeline (structured_agent's `app/` package):
                    INTENT graph  -> SEARCH + CLUSTER graph
  • chatbot_core/   the single-PDF Q&A chatbot (Qa.py + vectorizeer.py)

NOVA is the product. SONIC is the assistant persona that talks you through it.

Flow:
  USER RESEARCH IDEA
     -> INTENT agent frames it (Problem / Objective / Additional Context)
     -> you review/edit it
     -> SEARCH agent fetches + reranks + clusters papers
     -> results shown as clean thumbnails (title, authors, links)
     -> "Chat it out" on any paper: its PDF is downloaded, vectorized by
        vectorizeer.build_vectorstore, and you Q&A over it with Qa.py's chain.

This file is UI + wiring only. It does NOT change any agent or chatbot logic —
it imports their functions and drives them.

Where Streamlit re-executed one script top-to-bottom on every interaction, Gradio
builds a persistent component graph once and fires explicit handlers. So the
"stage" that Streamlit kept in session_state and branched on is here a set of
Columns whose `visible` flag every handler returns. Same state machine — declared
once instead of re-derived per rerun.

Run:
    python nova_app.py            # from inside the NOVA/ folder
"""

from ui import paths  # noqa: F401  — MUST be first: wires sys.path + chdir + .env

import shutil
import uuid

import gradio as gr

from ui import agents
from ui.agents import load_agents, preload_all
from ui.chat_engine import prepare_chat_stream
from ui.constants import CLUSTER_ACCENTS, SOURCE_COLORS
from ui.gpu import cuda_state
from ui.intent_text import join_intent_sections, split_intent_sections
from ui.papers import author_line, first_available
from ui.paths import DOWNLOADS_DIR, VECTORSTORES_DIR
from ui.search_progress import search_steps_html
from ui.sonic import SONIC_AVATAR, SONIC_DATA_URI, USER_AVATAR, sonic_says
from ui.theme import CSS, FORCE_DARK, NOVA_THEME

# ---------------------------------------------------------------------------
# 0. LOAD THE MODELS — HERE, AT IMPORT, ON THE MAIN THREAD.
#
# This placement is load-bearing, not stylistic. ZeroGPU forks its GPU worker
# from this process and `spaces`' torch patching (the thing that stops CUDA
# from really initialising before that fork) is thread-local to whichever
# thread called patch() — the main thread, during `import spaces`. Loading
# models from a Gradio worker thread or a daemon thread of our own escapes that
# patching and poisons the fork, which surfaces much later and very
# confusingly as:
#
#     RuntimeError: No CUDA GPUs are available   (in spaces' worker_init)
#
# So: main thread, before Gradio exists. See ui/gpu.py and ui/agents.py.
# ---------------------------------------------------------------------------
print(cuda_state("pre-preload"), flush=True)
preload_all()
print(cuda_state("post-preload"), flush=True)

# ---------------------------------------------------------------------------
# 1. SESSION STATE
# ---------------------------------------------------------------------------
STAGE_NAMES = ("boot", "welcome", "refining", "review", "searching", "results", "chat")


def new_state() -> dict:
    """One of these per browser session. Gradio deep-copies it into each new
    session, so the mutable members below are never shared across users."""
    return {
        "run_id": "",
        "user_query": "",
        "problem": "",
        "objective": "",
        "context": "",
        "clusters": [],
        "papers_by_key": {},   # normalized_title -> full record (flattened, for chat lookup)
        "active_chat": None,   # normalized_title of the paper being chatted, or None
        "chats": {},           # normalized_title -> {retriever, chain, messages, title}
        "source_status": {},   # {"Semantic Scholar": {"state": "rate_limited", ...}, ...}
    }


def wipe_disk_cache():
    """Delete every cached vectorstore and downloaded PDF so a new search starts
    from a clean slate — no old paper's chunks or PDFs can leak in."""
    for folder in (VECTORSTORES_DIR, DOWNLOADS_DIR):
        try:
            if folder.exists():
                shutil.rmtree(folder, ignore_errors=True)
            folder.mkdir(exist_ok=True)
        except Exception:
            pass


def _stages(active: str):
    """Visibility updates for every stage Column, in STAGE_NAMES order."""
    return tuple(gr.update(visible=(name == active)) for name in STAGE_NAMES)


# ---------------------------------------------------------------------------
# 2. STATIC MARKUP
# ---------------------------------------------------------------------------
HEADER_HTML = """
<div class="nova-brand">
  <span class="nova-star">✦</span>
  <span class="nova-mark">NOVA</span>
  <span class="nova-sub">Research&nbsp;Assistant</span>
  <span class="sonic-chip"><span class="sonic-dot"></span> SONIC online</span>
</div>
"""

HERO_FIGURE_HTML = (
    f'<div class="hero-figure"><img src="{SONIC_DATA_URI}" alt="SONIC"/>'
    f'<div class="hero-name">SONIC · your research buddy</div></div>'
)

HERO_SPEECH_HTML = """
<div class="hero-speech">
  <div class="sonic-name">SONIC</div>
  hey, wass up 👋<br>what's on your mind about research today?<br>
  Dump the raw idea on me — the messier the better. I'll shape it into something sharp.
</div>
<div class="hero-answer-label">✍️ your answer</div>
"""


def boot_html(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;{phase}&rdquo;</div>'
        f'  <div class="vec-bar"><div class="vec-fill" style="width:{pct}%"></div></div>'
        f'</div>'
    )


def source_status_html(status: dict) -> str:
    """Show per-source status so a rate-limited/failed source is never invisible."""
    if not status:
        return ""
    chips = []
    for name, s in status.items():
        state = (s or {}).get("state")
        if state == "ok":
            chips.append(f'<span class="src-stat ok">{name} ✓ {s.get("count", 0)}</span>')
        elif state == "rate_limited":
            chips.append(f'<span class="src-stat warn">{name} ⚠ rate-limited (HTTP {s.get("http", 429)})</span>')
        elif state == "error":
            detail = s.get("detail") or f'HTTP {s.get("http", "?")}'
            chips.append(f'<span class="src-stat err">{name} ✕ {detail}</span>')
        else:
            chips.append(f'<span class="src-stat muted">{name} —</span>')
    return '<div class="src-stat-row">' + "".join(chips) + "</div>"


def paper_card_html(norm_title: str, record: dict) -> str:
    title = record.get("title") or norm_title.title()
    badges = "".join(
        f'<span class="src-badge" style="color:{SOURCE_COLORS.get(s, "#8b93a7")};'
        f'border-color:{SOURCE_COLORS.get(s, "#8b93a7")}55;'
        f'background:{SOURCE_COLORS.get(s, "#8b93a7")}18;">{s}</span>'
        for s in (record.get("source") or [])
    )
    return (
        f'<div class="paper-title">{title}</div>'
        f'<div class="paper-meta">{author_line(record.get("authors"), record.get("year"))}<br>{badges}</div>'
    )


def card_links_html(page_url: str, pdf_url: str) -> str:
    """The 📄 Paper / ⬇ PDF pair. Plain anchors rather than gr.Button: they're
    pure navigation, and a real <a> opens a new tab with no server round-trip."""
    paper = (f'<a class="card-link" href="{page_url}" target="_blank" rel="noopener">📄 Paper</a>'
             if page_url else '<span class="card-link dead">📄 Paper</span>')
    pdf = (f'<a class="card-link" href="{pdf_url}" target="_blank" rel="noopener">⬇ PDF</a>'
           if pdf_url else '<span class="card-link dead">⬇ PDF</span>')
    return f'<div class="card-links">{paper}{pdf}</div>'


# ---------------------------------------------------------------------------
# 3. HANDLERS THAT TOUCH NO COMPONENTS
# ---------------------------------------------------------------------------
def open_chat_for(norm_title: str):
    """Build a per-card click handler. The card grid is generated in a loop, so
    each button needs to close over its own paper key."""
    def _open(st):
        st["active_chat"] = norm_title
        record = st["papers_by_key"].get(norm_title, {})
        title = record.get("title") or norm_title.title()
        head = (f'<div class="cluster-head"><div class="cluster-bar" style="background:#7c5cff;"></div>'
                f'<div class="cluster-title">💬 {title}</div></div>')
        cached = st["chats"].get(norm_title)
        return (*_stages("chat"), head,
                gr.update(value="", visible=not cached),
                gr.update(value=(cached["messages"] if cached else []), visible=bool(cached)),
                gr.update(visible=bool(cached)),
                st)
    return _open


def prep_chat(st):
    """Download + vectorize this paper, animating SONIC's pep-quotes while the
    real work runs on a worker thread. No-op if this chat is already built."""
    key = st["active_chat"]
    if not key or key in st["chats"]:
        return

    record = st["papers_by_key"].get(key, {})
    session = error = None
    for html, done, session, error in prepare_chat_stream(record):
        if not done:
            yield (gr.update(value=html, visible=True), gr.update(visible=False),
                   gr.update(visible=False), st)

    if error or not session:
        msg = error or "Couldn't prepare this paper for chat."
        yield (gr.update(value=f'<div class="nova-error">{msg}</div>', visible=True),
               gr.update(visible=False), gr.update(visible=False), st)
        return

    session.update({"messages": [], "title": record.get("title") or key.title()})
    st["chats"][key] = session
    yield (gr.update(value="", visible=False), gr.update(value=[], visible=True),
           gr.update(visible=True), st)


# ---------------------------------------------------------------------------
# 4. THE APP
# ---------------------------------------------------------------------------
# Gradio 6 moved theme/css/js off the Blocks constructor and onto launch().
with gr.Blocks(title="NOVA · Research Assistant", analytics_enabled=False) as demo:
    state = gr.State(new_state())
    # Mirrors state["clusters"]. gr.render can't watch a dict mutated in place,
    # so the search handler reassigns this to a fresh list to trigger a redraw.
    clusters_state = gr.State([])

    with gr.Column(elem_id="nova-root"):
        gr.HTML(HEADER_HTML)

        # ---------------- BOOT ----------------
        with gr.Column(visible=True) as boot_col:
            boot_panel = gr.HTML(boot_html("Waking up SONIC — loading the research + reading models…", 8))

        # ---------------- WELCOME ----------------
        with gr.Column(visible=False) as welcome_col:
            with gr.Row(equal_height=False):
                with gr.Column(scale=9):
                    gr.HTML(HERO_FIGURE_HTML)
                with gr.Column(scale=11):
                    gr.HTML(HERO_SPEECH_HTML)
                    query_box = gr.Textbox(
                        lines=6, max_lines=12, show_label=False, container=False,
                        placeholder="e.g. I want to compare fuel efficiency of human-driven vs RL-controlled "
                                    "cars in car-following… comparing is hard because velocity, acceleration, "
                                    "headway all change at once…",
                    )
                    go_btn = gr.Button("Let's go  ✦", variant="primary")

        # ---------------- REFINING ----------------
        with gr.Column(visible=False) as refining_col:
            gr.HTML(sonic_says("that seems great — lemme juss refine it ✨"))
            refining_panel = gr.HTML()

        # ---------------- REVIEW ----------------
        with gr.Column(visible=False) as review_col:
            gr.HTML(sonic_says("here's how I framed it. Tweak anything that's off, then I'll go hunting 🔍"))
            gr.HTML('<div class="field-label">🧩 Problem</div>')
            problem_box = gr.Textbox(lines=5, show_label=False, container=False)
            gr.HTML('<div class="field-label">🎯 Objective</div>')
            objective_box = gr.Textbox(lines=4, show_label=False, container=False)
            gr.HTML('<div class="field-label">🗂️ Additional Context</div>')
            context_box = gr.Textbox(lines=4, show_label=False, container=False)
            with gr.Row():
                find_btn = gr.Button("Find the papers  🔍", variant="primary", scale=2)
                over_btn = gr.Button("Start over", scale=1)
                gr.HTML("")   # spacer: keeps the two buttons off full width

        # ---------------- SEARCHING ----------------
        with gr.Column(visible=False) as searching_col:
            gr.HTML(sonic_says("on it — scouring arXiv, Semantic Scholar &amp; OpenAlex, then reranking "
                               "and clustering by approach 🔎"))
            search_panel = gr.HTML()

        # ---------------- RESULTS ----------------
        # Body is filled in by the @gr.render below, once every component it
        # needs to drive (the chat stage) exists.
        with gr.Column(visible=False) as results_col:
            results_head = gr.HTML()
            with gr.Row():
                new_search_btn = gr.Button("🔄 New search", scale=1)
                gr.HTML("")   # spacer

        # ---------------- CHAT ----------------
        with gr.Column(visible=False) as chat_col:
            with gr.Row():
                back_btn = gr.Button("← Back to papers", scale=1)
                gr.HTML("")   # spacer
            chat_title = gr.HTML()
            chat_vec = gr.HTML()
            # Gradio 6 speaks the {"role","content"} message format natively —
            # no type="messages" to opt into it any more.
            chatbot = gr.Chatbot(
                height=520, show_label=False, visible=False,
                elem_id="nova-chat", avatar_images=(USER_AVATAR, SONIC_AVATAR),
                placeholder="ask me anything about this paper — I've read every page 📄",
            )
            with gr.Row(visible=False) as chat_input_row:
                chat_input = gr.Textbox(show_label=False, container=False, scale=9,
                                        placeholder="Ask about this paper…")
                send_btn = gr.Button("Send", variant="primary", scale=1)

    STAGE_COLS = [boot_col, welcome_col, refining_col, review_col, searching_col, results_col, chat_col]

    # Re-enter the results Column now that the chat components exist, so each
    # card's "Chat it out" button can wire straight into them.
    with results_col:
        @gr.render(inputs=[clusters_state, state], triggers=[clusters_state.change])
        def draw_results(clusters, st):
            """Redrawn whenever a search completes. Streamlit rebuilt this grid on
            every rerun for free; in Gradio the per-card buttons need real event
            handlers, so the whole thing is (re)declared here."""
            if not clusters:
                return
            for i, cluster in enumerate(clusters):
                papers = cluster.get("papers") or {}
                if not papers:
                    continue
                accent = CLUSTER_ACCENTS[i % len(CLUSTER_ACCENTS)]
                gr.HTML(
                    f'<div class="cluster-head">'
                    f'  <div class="cluster-bar" style="background:{accent};"></div>'
                    f'  <div class="cluster-title">{cluster.get("label", "Approach")}</div>'
                    f'</div>'
                )
                if cluster.get("rationale"):
                    gr.HTML(f'<div class="cluster-why">{cluster["rationale"]}</div>')

                items = list(papers.items())
                for row_start in range(0, len(items), 2):
                    with gr.Row(equal_height=True):
                        for norm_title, record in items[row_start:row_start + 2]:
                            with gr.Column(elem_classes=["paper-card"]):
                                gr.HTML(paper_card_html(norm_title, record))
                                page_url = first_available(record.get("url"))
                                pdf_url = first_available(record.get("pdf_url"))
                                gr.HTML(card_links_html(page_url, pdf_url))
                                # The real PDF is resolved/verified on click (the
                                # HYBRID deep step), so any link is enough to try.
                                chat_btn = gr.Button(
                                    "💬 Chat it out", variant="primary", size="sm",
                                    interactive=bool(page_url or pdf_url),
                                )
                                chat_btn.click(
                                    open_chat_for(norm_title),
                                    inputs=[state],
                                    outputs=[*STAGE_COLS, chat_title, chat_vec, chatbot,
                                             chat_input_row, state],
                                ).then(
                                    prep_chat,
                                    inputs=[state],
                                    outputs=[chat_vec, chatbot, chat_input_row, state],
                                )

    # -----------------------------------------------------------------------
    # 5. WIRING
    # -----------------------------------------------------------------------
    def do_boot():
        """Runs once per page load. Almost nothing left to do.

        Every model is already resident: preload_all() ran at import, on the
        main thread, because ZeroGPU requires it (see section 0). So this is now
        just the splash -> welcome transition, plus surfacing a load failure that
        preload_all() deliberately swallowed rather than killing the Space with.
        """
        if agents.BOOT_ERROR:
            yield (*_stages("boot"),
                   f'<div class="nova-error">SONIC couldn\'t wake up: {agents.BOOT_ERROR}<br>'
                   f'Check that GROQ_API_KEY / SECOND_GROQ_API_KEY / TAVILY_API_KEY are set.</div>')
            return
        yield (*_stages("welcome"), "")

    demo.load(do_boot, outputs=[*STAGE_COLS, boot_panel])

    def go(query, st):
        """WELCOME -> REFINING -> REVIEW. Runs the INTENT graph up to its
        human-review interrupt, then hands the framed sections to the form."""
        if not (query or "").strip():
            gr.Warning("Give me something to work with first 🙂")
            yield (*_stages("welcome"), gr.update(), gr.update(), gr.update(), st)
            return

        st["user_query"] = query.strip()
        st["run_id"] = str(uuid.uuid4())
        yield (*_stages("refining"), gr.update(), gr.update(), gr.update(), st)

        intent_graph, _ = load_agents()
        config = {"configurable": {"thread_id": st["run_id"]}}
        try:
            result = intent_graph.invoke({"user_query": st["user_query"], "run_id": st["run_id"]},
                                         config=config)
            payload = result["__interrupt__"][0].value
            problem, objective, context = split_intent_sections(payload["polished_research_intent"])
        except Exception as e:
            gr.Warning(f"Something went sideways: {type(e).__name__}: {e}")
            yield (*_stages("welcome"), gr.update(), gr.update(), gr.update(), st)
            return

        st["problem"], st["objective"], st["context"] = problem, objective, context
        yield (*_stages("review"), problem, objective, context, st)

    go_btn.click(go, inputs=[query_box, state],
                 outputs=[*STAGE_COLS, problem_box, objective_box, context_box, state])

    def find(problem, objective, context, st):
        """REVIEW -> SEARCHING -> RESULTS. Resumes the INTENT graph past its
        interrupt, then STREAMS the SEARCH + CLUSTER graph so the checklist ticks
        each step off live instead of hanging on one spinner."""
        from langgraph.types import Command

        if not (problem or "").strip() or not (objective or "").strip():
            gr.Warning("Problem and Objective can't be empty.")
            yield (*_stages("review"), gr.update(), gr.update(), st, gr.update())
            return

        st["problem"], st["objective"], st["context"] = problem, objective, context
        yield (*_stages("searching"), search_steps_html(set(), {}), gr.update(), st, gr.update())

        intent_graph, search_graph = load_agents()
        config = {"configurable": {"thread_id": st["run_id"]}}
        edited_intent = join_intent_sections(problem, objective, context)

        try:
            resume_result = intent_graph.invoke(Command(resume=edited_intent), config=config)
            human_verified_intent = resume_result["human_verified_intent"]

            completed, counts, final_state = set(), {}, {}
            source_field = {"arxiv": "arXiv_paper", "semantic_scholar": "Semantic_Scholar_paper",
                            "open_alex": "Open_Alex_paper"}
            for update in search_graph.stream(
                {"ResearchIntent": human_verified_intent, "run_id": st["run_id"]},
                stream_mode="updates",
            ):
                for node_name, delta in update.items():
                    completed.add(node_name)
                    if isinstance(delta, dict):
                        final_state.update(delta)
                        if node_name in source_field:
                            counts[node_name] = len(delta.get(source_field[node_name]) or [])
                yield (*_stages("searching"), search_steps_html(completed, counts),
                       gr.update(), st, gr.update())
        except Exception as e:
            gr.Warning(f"Something went sideways: {type(e).__name__}: {e}")
            yield (*_stages("review"), gr.update(), gr.update(), st, gr.update())
            return

        clusters = final_state.get("clustered_papers") or []
        # flatten every paper into a lookup keyed by its normalized_title (the
        # cluster dict key) so "Chat it out" can find the record anywhere.
        papers_by_key = {}
        for cluster in clusters:
            for norm_title, record in (cluster.get("papers") or {}).items():
                papers_by_key[norm_title] = record

        st["clusters"] = clusters
        st["papers_by_key"] = papers_by_key
        st["source_status"] = {
            "arXiv": final_state.get("arxiv_status") or {},
            "Semantic Scholar": final_state.get("semantic_scholar_status") or {},
            "OpenAlex": final_state.get("open_alex_status") or {},
        }

        total = sum(len(c.get("papers") or {}) for c in clusters)
        if total == 0:
            head = sonic_says("hmm, I couldn't pull solid matches for that one — see the source status below. "
                              "If a source is rate-limited, that's usually why. Try again in a bit, or "
                              "loosen the framing.")
        else:
            head = sonic_says(f"these are the best matches 🎯<br>{total} papers, grouped into "
                              f"{len(clusters)} approaches. Hit <b>Chat it out</b> on any paper to "
                              f"actually talk to it.")
        head += source_status_html(st["source_status"])

        yield (*_stages("results"), gr.update(), head, st, list(clusters))

    find_btn.click(find, inputs=[problem_box, objective_box, context_box, state],
                   outputs=[*STAGE_COLS, search_panel, results_head, state, clusters_state])

    def start_over():
        """Full memory refresh: clear this session's results + open chats, and wipe
        the on-disk vectorstore/PDF caches too."""
        wipe_disk_cache()
        return (*_stages("welcome"), "", new_state(), [])

    over_btn.click(start_over, outputs=[*STAGE_COLS, query_box, state, clusters_state])
    new_search_btn.click(start_over, outputs=[*STAGE_COLS, query_box, state, clusters_state])

    def answer(message, st):
        """Stream one grounded answer, then append the page citations."""
        from langchain_core.messages import AIMessage, HumanMessage
        from Qa import format_docs

        message = (message or "").strip()
        if not message or not st.get("active_chat"):
            yield gr.update(), ""
            return

        session = st["chats"][st["active_chat"]]
        session["messages"].append({"role": "user", "content": message})
        yield [dict(m) for m in session["messages"]], ""

        # LangChain chat history from prior turns (excludes the just-added question)
        history = [HumanMessage(content=m["content"]) if m["role"] == "user"
                   else AIMessage(content=m["content"])
                   for m in session["messages"][:-1]]

        try:
            docs = session["retriever"].invoke(message)
            context = format_docs(docs)
            session["messages"].append({"role": "assistant", "content": ""})
            for chunk in session["chain"].stream(
                {"question": message, "chat_history": history, "context": context}
            ):
                if chunk.content:
                    session["messages"][-1]["content"] += chunk.content
                    yield [dict(m) for m in session["messages"]], ""
            pages = sorted({f"p.{d.metadata.get('page')}" for d in docs})
            if pages:
                session["messages"][-1]["content"] += "\n\n*sources: " + ", ".join(pages) + "*"
        except Exception as e:
            session["messages"].append(
                {"role": "assistant",
                 "content": f"Sorry — I hit an error answering that: {type(e).__name__}: {e}"}
            )
        yield [dict(m) for m in session["messages"]], ""

    for trigger in (chat_input.submit, send_btn.click):
        trigger(answer, inputs=[chat_input, state], outputs=[chatbot, chat_input])

    def back_to_papers(st):
        st["active_chat"] = None
        return (*_stages("results"), st)

    back_btn.click(back_to_papers, inputs=[state], outputs=[*STAGE_COLS, state])


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=4).launch(
        theme=NOVA_THEME, css=CSS, js=FORCE_DARK,
        server_name="0.0.0.0", server_port=7860,
    )