app.py

import os
import gradio as gr
import fitz
import math, re, json, time, html
from google import genai
from google.genai import types

GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")

def _with_retry(fn, max_retries=4):
    """Call fn(), retrying on 429 with exponential backoff."""
    for attempt in range(max_retries):
        try:
            return fn()
        except Exception as e:
            msg = str(e)
            if "429" in msg or "RESOURCE_EXHAUSTED" in msg:
                # Parse retryDelay from error if present, else back off exponentially
                wait = 10 * (2 ** attempt)
                import re as _re
                m = _re.search(r"retryDelay.*?(\d+)s", msg)
                if m:
                    wait = int(m.group(1)) + 2
                time.sleep(wait)
            else:
                raise
    raise RuntimeError(f"Failed after {max_retries} retries due to rate limits. Try again in a minute.")

# ── Chunking ──────────────────────────────────────────────────────────────────

CHUNK_SIZE    = 800
CHUNK_OVERLAP = 120

def split_into_chunks(text: str) -> list:
    paras  = [p.strip() for p in re.split(r"\n{2,}", text) if p.strip()]
    chunks = []
    for para in paras:
        if len(para) <= CHUNK_SIZE:
            chunks.append(para)
        else:
            i = 0
            while i < len(para):
                chunks.append(para[i : i + CHUNK_SIZE].strip())
                if i + CHUNK_SIZE >= len(para):
                    break
                i += CHUNK_SIZE - CHUNK_OVERLAP
    return [c for c in chunks if c]

# ── Dense vector store ────────────────────────────────────────────────────────

def cosine(a, b):
    dot = sum(x*y for x,y in zip(a,b))
    na  = math.sqrt(sum(x*x for x in a))
    nb  = math.sqrt(sum(x*x for x in b))
    return dot/(na*nb) if na and nb else 0.0

class EmbeddingStore:
    MODEL      = "gemini-embedding-2-preview"
    DIM        = 1536
    BATCH_SIZE = 20

    def __init__(self):
        self.chunks   = []
        self.client   = None
        self._key     = ""

    def set_client(self, key):
        if key != self._key:
            self.client = genai.Client(api_key=key)
            self._key   = key

    def clear(self):
        self.chunks = []

    def _embed(self, texts, task):
        r = self.client.models.embed_content(
            model=self.MODEL, contents=texts,
            config=types.EmbedContentConfig(task_type=task, output_dimensionality=self.DIM),
        )
        return [list(e.values) for e in r.embeddings]

    def add_chunks(self, texts):
        added = 0
        for i in range(0, len(texts), self.BATCH_SIZE):
            batch  = texts[i : i + self.BATCH_SIZE]
            embeds = self._embed(batch, "RETRIEVAL_DOCUMENT")
            for t, e in zip(batch, embeds):
                self.chunks.append({"id": f"chunk_{len(self.chunks)}", "text": t, "emb": e})
                added += 1
            if i + self.BATCH_SIZE < len(texts):
                time.sleep(0.3)
        return added

    def search(self, query, k=5):
        if not self.chunks: return []
        qvec = self._embed(query, "RETRIEVAL_QUERY")[0]
        sc   = [{"id":c["id"],"text":c["text"],"score":cosine(qvec,c["emb"])} for c in self.chunks]
        sc.sort(key=lambda x: x["score"], reverse=True)
        return sc[:k]

    def get(self, cid):
        return next((c for c in self.chunks if c["id"]==cid), None)

    def overview(self):
        return [{"id":c["id"],"preview":c["text"][:120]+"..."} for c in self.chunks]

    @property
    def size(self): return len(self.chunks)

# ── PDF extraction ─────────────────────────────────────────────────────────────

def extract_pdf(path):
    doc  = fitz.open(path)
    text = "\n\n".join(p.get_text() for p in doc)
    doc.close()
    return text

# ── HTML rendering helpers ────────────────────────────────────────────────────
# These build the live agent visualization panel

def esc(s): return html.escape(str(s))

PANEL_CSS = """
<style>
@import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:ital,wght@0,400;0,500;1,400&family=IBM+Plex+Sans:wght@300;400;500;600&display=swap');

.av-root {
  font-family: 'IBM Plex Sans', sans-serif;
  font-size: 13px;
  line-height: 1.55;
  color: #1a1a2e;
  display: flex;
  flex-direction: column;
  gap: 0;
  background: #f8f9fc;
  border-radius: 10px;
  overflow: hidden;
  border: 1px solid #e2e5f0;
  min-height: 80px;
}

/* Step cards */
.av-step {
  border-left: 3px solid #e2e5f0;
  margin: 0;
  padding: 12px 16px;
  background: #fff;
  border-bottom: 1px solid #f0f2f8;
  animation: avSlide 0.25s ease;
  position: relative;
}
@keyframes avSlide {
  from { opacity: 0; transform: translateY(-6px); }
  to   { opacity: 1; transform: none; }
}

/* Color-coded left border per type */
.av-step.type-think    { border-left-color: #7c3aed; background: #faf5ff; }
.av-step.type-search   { border-left-color: #0891b2; background: #f0f9ff; }
.av-step.type-fetch    { border-left-color: #059669; background: #f0fdf4; }
.av-step.type-list     { border-left-color: #d97706; background: #fffbeb; }
.av-step.type-results  { border-left-color: #0891b2; background: #f8fdff; }
.av-step.type-answer   { border-left-color: #16a34a; background: #f0fdf4; }
.av-step.type-error    { border-left-color: #dc2626; background: #fef2f2; }
.av-step.type-embed    { border-left-color: #7c3aed; background: #faf5ff; }

/* Step header row */
.av-header {
  display: flex;
  align-items: center;
  gap: 8px;
  margin-bottom: 4px;
}
.av-icon {
  font-size: 14px;
  flex-shrink: 0;
  line-height: 1;
}
.av-tag {
  font-family: 'IBM Plex Mono', monospace;
  font-size: 10px;
  font-weight: 500;
  text-transform: uppercase;
  letter-spacing: 0.08em;
  padding: 2px 8px;
  border-radius: 20px;
  flex-shrink: 0;
}
.type-think  .av-tag { background: #ede9fe; color: #5b21b6; }
.type-search .av-tag { background: #e0f2fe; color: #0369a1; }
.type-fetch  .av-tag { background: #dcfce7; color: #166534; }
.type-list   .av-tag { background: #fef3c7; color: #92400e; }
.type-results .av-tag { background: #e0f2fe; color: #0369a1; }
.type-answer  .av-tag { background: #dcfce7; color: #166534; }
.type-error   .av-tag { background: #fee2e2; color: #991b1b; }
.type-embed   .av-tag { background: #ede9fe; color: #5b21b6; }

.av-step-num {
  font-family: 'IBM Plex Mono', monospace;
  font-size: 10px;
  color: #94a3b8;
  margin-left: auto;
}

/* Body text */
.av-body {
  font-size: 12.5px;
  color: #374151;
  line-height: 1.6;
}
.av-body.mono {
  font-family: 'IBM Plex Mono', monospace;
  font-size: 11.5px;
  color: #1e293b;
}

/* Query line */
.av-query {
  font-family: 'IBM Plex Mono', monospace;
  font-size: 12px;
  color: #0369a1;
  background: #e0f2fe;
  padding: 4px 10px;
  border-radius: 4px;
  display: inline-block;
  margin-top: 3px;
  word-break: break-word;
}

/* Chunk cards inside results */
.av-chunks { display: flex; flex-direction: column; gap: 6px; margin-top: 8px; }
.av-chunk {
  background: #fff;
  border: 1px solid #bae6fd;
  border-radius: 6px;
  padding: 8px 11px;
}
.av-chunk-header {
  display: flex;
  align-items: center;
  gap: 8px;
  margin-bottom: 4px;
}
.av-chunk-id {
  font-family: 'IBM Plex Mono', monospace;
  font-size: 11px;
  font-weight: 500;
  color: #0369a1;
  background: #e0f2fe;
  padding: 1px 7px;
  border-radius: 20px;
}
.av-score-bar {
  flex: 1;
  height: 5px;
  background: #e2e8f0;
  border-radius: 3px;
  overflow: hidden;
}
.av-score-fill {
  height: 100%;
  background: linear-gradient(90deg, #38bdf8, #0369a1);
  border-radius: 3px;
  transition: width 0.5s ease;
}
.av-score-val {
  font-family: 'IBM Plex Mono', monospace;
  font-size: 10px;
  color: #64748b;
  white-space: nowrap;
}
.av-chunk-text {
  font-size: 11.5px;
  color: #475569;
  line-height: 1.5;
  display: -webkit-box;
  -webkit-line-clamp: 3;
  -webkit-box-orient: vertical;
  overflow: hidden;
}

/* Divider between agent runs */
.av-divider {
  text-align: center;
  padding: 8px 0;
  font-family: 'IBM Plex Mono', monospace;
  font-size: 10px;
  color: #94a3b8;
  letter-spacing: 0.1em;
  background: #f8f9fc;
  border-bottom: 1px solid #e2e5f0;
}

/* Empty state */
.av-empty {
  display: flex;
  flex-direction: column;
  align-items: center;
  justify-content: center;
  padding: 40px 20px;
  gap: 8px;
  color: #94a3b8;
  font-family: 'IBM Plex Mono', monospace;
  font-size: 12px;
  text-align: center;
  min-height: 120px;
}
.av-empty-icon { font-size: 28px; }

/* Spinning indicator */
.av-thinking {
  display: inline-block;
  width: 10px; height: 10px;
  border: 2px solid #c4b5fd;
  border-top-color: #7c3aed;
  border-radius: 50%;
  animation: avSpin 0.7s linear infinite;
  vertical-align: middle;
  margin-right: 5px;
}
@keyframes avSpin { to { transform: rotate(360deg); } }
</style>
"""

EMPTY_PANEL = PANEL_CSS + """
<div class="av-root">
  <div class="av-empty">
    <div>Agent decisions will appear here in real-time</div>
    <div style="color:#cbd5e1; font-size:11px; margin-top:2px">Upload a PDF → ask a question</div>
  </div>
</div>
"""

def render_panel(steps: list) -> str:
    """Render the full agent visualization panel from a list of step dicts."""
    if not steps:
        return EMPTY_PANEL

    cards = []
    for i, s in enumerate(steps):
        t    = s["type"]
        body = ""

        if t == "divider":
            cards.append(f'<div class="av-divider">── {esc(s.get("label",""))} ──</div>')
            continue

        elif t == "think":
            body = f'<div class="av-body">{esc(s["text"])}</div>'

        elif t == "search":
            body = (
                f'<div class="av-body">Querying corpus for:</div>'
                f'<div class="av-query">{esc(s["query"])}</div>'
                + (f'<div class="av-body" style="margin-top:6px;color:#64748b">k = {s["k"]}</div>' if s.get("k") else "")
            )

        elif t == "results":
            results = s.get("results", [])
            chunk_html = ""
            for r in results:
                score   = r["score"]
                bar_pct = min(int(score * 100 * 2.5), 100)  # scale cosine to visual
                preview = r["text"][:240] + ("…" if len(r["text"]) > 240 else "")
                chunk_html += f"""
                <div class="av-chunk">
                  <div class="av-chunk-header">
                    <span class="av-chunk-id">{esc(r["id"])}</span>
                    <div class="av-score-bar"><div class="av-score-fill" style="width:{bar_pct}%"></div></div>
                    <span class="av-score-val">sim {score:.3f}</span>
                  </div>
                  <div class="av-chunk-text">{esc(preview)}</div>
                </div>"""
            body = (
                f'<div class="av-body" style="color:#0369a1;font-weight:500">'
                f'Found {len(results)} relevant chunk{"s" if len(results)!=1 else ""}</div>'
                f'<div class="av-chunks">{chunk_html}</div>'
            )

        elif t == "fetch":
            body = (
                f'<div class="av-body">Fetching full text of:</div>'
                f'<div class="av-query">{esc(s["chunk_id"])}</div>'
            )
            if s.get("text"):
                preview = s["text"][:300] + ("…" if len(s["text"]) > 300 else "")
                body += f'<div class="av-body mono" style="margin-top:8px;padding:8px;background:#f0fdf4;border-radius:5px;border:1px solid #bbf7d0">{esc(preview)}</div>'

        elif t == "list":
            body = f'<div class="av-body">Scanning all {s.get("total",0)} chunks in corpus for context…</div>'

        elif t == "answer":
            body = (
                f'<div class="av-body" style="color:#166534;font-weight:500">Answer synthesized</div>'
                f'<div class="av-body" style="margin-top:4px;color:#64748b">from {s.get("chunks_used",0)} retrieved chunk(s) across {s.get("steps",0)} retrieval step(s)</div>'
            )

        elif t == "error":
            body = f'<div class="av-body" style="color:#dc2626">{esc(s["text"])}</div>'

        elif t == "embed":
            body = f'<div class="av-body">{esc(s["text"])}</div>'

        type_icons = {
            "think":   ("💭", "Thinking"),
            "search":  ("🔍", "Search"),
            "results": ("📋", "Retrieved"),
            "fetch":   ("📄", "Fetch chunk"),
            "list":    ("📚", "List all"),
            "answer":  ("✅", "Answer ready"),
            "error":   ("❌", "Error"),
            "embed":   ("⚡", "Embedding"),
        }
        icon, tag = type_icons.get(t, ("·", t))
        step_num  = s.get("step_num", "")

        cards.append(f"""
        <div class="av-step type-{t}">
          <div class="av-header">
            <span class="av-icon">{icon}</span>
            <span class="av-tag">{tag}</span>
            {"<span class='av-step-num'>step "+str(step_num)+"</span>" if step_num else ""}
          </div>
          {body}
        </div>""")

    return PANEL_CSS + '<div class="av-root">' + "".join(cards) + "</div>"

# ── Agent tools ───────────────────────────────────────────────────────────────

TOOL_DECLARATIONS = [
    types.FunctionDeclaration(
        name="search_chunks",
        description=(
            "Semantically search the PDF corpus using Gemini Embedding 2 dense vectors. "
            "Returns top-k chunks ranked by cosine similarity. "
            "Call multiple times with different queries for multi-part questions."
        ),
        parameters=types.Schema(
            type=types.Type.OBJECT,
            properties={
                "query": types.Schema(type=types.Type.STRING, description="Focused natural-language search query"),
                "k":     types.Schema(type=types.Type.INTEGER, description="Number of chunks to return (1-8, default 4)"),
            },
            required=["query"],
        ),
    ),
    types.FunctionDeclaration(
        name="get_chunk_by_id",
        description="Retrieve full text of a specific chunk by ID. Use when a search preview isn't enough.",
        parameters=types.Schema(
            type=types.Type.OBJECT,
            properties={"chunk_id": types.Schema(type=types.Type.STRING, description="e.g. chunk_5")},
            required=["chunk_id"],
        ),
    ),
    types.FunctionDeclaration(
        name="list_all_chunks",
        description="Get overview of all chunks (IDs + previews). Use to understand corpus scope before searching.",
        parameters=types.Schema(type=types.Type.OBJECT, properties={}),
    ),
]

AGENT_TOOLS   = [types.Tool(function_declarations=TOOL_DECLARATIONS)]

SYSTEM_PROMPT = """You are a precise RAG agent with access to a PDF corpus via semantic search tools.

To answer questions:
1. PLAN what you need
2. RETRIEVE via search_chunks — use multiple focused queries for multi-part questions
3. EXPAND with get_chunk_by_id for full chunk text when needed
4. SYNTHESIZE a final grounded answer

Rules:
- Always retrieve before answering — never rely on prior knowledge
- Cite chunk IDs inline e.g. [chunk_2] for every factual claim
- If chunks lack sufficient info, say so clearly
- Be thorough but concise"""

# ── Agentic loop ──────────────────────────────────────────────────────────────

def run_agent(query, store, history, steps):
    """Generator: yields (history, steps, panel_html) on every decision."""

    def emit(h, s):
        return h, s, render_panel(s)

    if not GEMINI_API_KEY:
        steps = steps + [{"type": "error", "text": "GEMINI_API_KEY secret not set in Space settings."}]
        yield emit(history, steps)
        return
    if not store or store.size == 0:
        steps = steps + [{"type": "error", "text": "No PDF loaded. Upload and embed first."}]
        yield emit(history, steps)
        return
    if not query.strip():
        return

    store.set_client(GEMINI_API_KEY)
    client = genai.Client(api_key=GEMINI_API_KEY)

    history = history + [{"role": "user", "content": query}]
    steps   = steps + [{"type": "divider", "label": query[:60] + ("…" if len(query) > 60 else "")}]
    yield emit(history, steps)

    gemini_msgs = []
    for msg in history:
        role = "user" if msg["role"] == "user" else "model"
        gemini_msgs.append(types.Content(role=role, parts=[types.Part(text=msg["content"])]))

    MAX_STEPS     = 12
    step          = 0
    retrieval_n   = 0
    chunks_used   = set()

    while step < MAX_STEPS:
        step += 1

        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=gemini_msgs,
            config=types.GenerateContentConfig(
                system_instruction=SYSTEM_PROMPT,
                tools=AGENT_TOOLS,
                temperature=0.2,
                max_output_tokens=2048,
            ),
        )

        candidate = response.candidates[0]
        parts     = candidate.content.parts
        gemini_msgs.append(types.Content(role="model", parts=parts))

        has_tool   = False
        tool_resps = []
        text_parts = []

        for part in parts:
            if hasattr(part, "text") and part.text and part.text.strip():
                txt = part.text.strip()
                # Show thinking only if it's not the final answer text
                steps = steps + [{"type": "think", "text": txt[:500] + ("…" if len(txt) > 500 else ""), "step_num": step}]
                yield emit(history, steps)
                text_parts.append(part.text)

            if hasattr(part, "function_call") and part.function_call:
                has_tool    = True
                fn          = part.function_call
                name        = fn.name
                args        = dict(fn.args) if fn.args else {}
                retrieval_n += 1

                if name == "search_chunks":
                    q = args.get("query", "")
                    k = int(args.get("k", 4))
                    steps = steps + [{"type": "search", "query": q, "k": k, "step_num": step}]
                    yield emit(history, steps)

                    results = store.search(q, k=k)
                    for r in results:
                        chunks_used.add(r["id"])

                    steps = steps + [{"type": "results", "results": results, "step_num": step}]
                    yield emit(history, steps)

                    result = {"results": [{"id":r["id"],"score":round(r["score"],4),"text":r["text"]} for r in results]}

                elif name == "get_chunk_by_id":
                    cid   = args.get("chunk_id", "")
                    chunk = store.get(cid)
                    steps = steps + [{"type": "fetch", "chunk_id": cid,
                                      "text": chunk["text"] if chunk else None, "step_num": step}]
                    yield emit(history, steps)
                    result = {"id": chunk["id"], "text": chunk["text"]} if chunk else {"error": "Not found"}
                    if chunk:
                        chunks_used.add(chunk["id"])

                elif name == "list_all_chunks":
                    steps = steps + [{"type": "list", "total": store.size, "step_num": step}]
                    yield emit(history, steps)
                    result = {"chunks": store.overview()}

                else:
                    result = {"error": f"Unknown tool: {name}"}

                tool_resps.append(types.Part(
                    function_response=types.FunctionResponse(name=name, response=result)
                ))

        if has_tool and tool_resps:
            gemini_msgs.append(types.Content(role="user", parts=tool_resps))
            continue

        # Final answer
        final = "\n".join(text_parts).strip() or "(No response)"
        history = history + [{"role": "assistant", "content": final}]
        steps   = steps + [{"type": "answer", "chunks_used": len(chunks_used),
                             "steps": retrieval_n, "step_num": step}]
        yield emit(history, steps)
        return

    steps = steps + [{"type": "error", "text": f"Reached max steps ({MAX_STEPS})."}]
    yield emit(history, steps)

# ── PDF processing ─────────────────────────────────────────────────────────────

def process_pdf(pdf_file, store, steps):
    if pdf_file is None:
        steps = steps + [{"type": "error", "text": "No file uploaded."}]
        yield store, "⚠️ No file uploaded.", "0 chunks", render_panel(steps), steps
        return
    if not GEMINI_API_KEY:
        steps = steps + [{"type": "error", "text": "GEMINI_API_KEY secret not set in Space settings."}]
        yield store, "❌ GEMINI_API_KEY not set.", "0 chunks", render_panel(steps), steps
        return

    try:
        new_store = EmbeddingStore()
        new_store.set_client(GEMINI_API_KEY)

        steps = steps + [{"type": "embed", "text": "📄 Extracting text from PDF…"}]
        yield new_store, "📄 Extracting text…", "…", render_panel(steps), steps

        text = extract_pdf(pdf_file.name)
        if not text.strip():
            steps = steps + [{"type": "error", "text": "No extractable text in PDF."}]
            yield new_store, "❌ No text found.", "0 chunks", render_panel(steps), steps
            return

        chars  = len(text)
        chunks = split_into_chunks(text)
        total  = len(chunks)

        steps = steps + [{"type": "embed", "text": f"✂️ Split into {total} chunks · embedding…"}]
        yield new_store, f"⚡ Embedding {total} chunks…", f"0 / {total}", render_panel(steps), steps

        added = new_store.add_chunks(chunks)

        steps = steps + [{"type": "embed", "text": f"✅ {added} chunks ready"}]
        msg   = f"✅ Ready — {added} chunks · {chars:,} chars"
        yield new_store, msg, f"{added} chunks", render_panel(steps), steps

    except Exception as e:
        steps = steps + [{"type": "error", "text": str(e)}]
        yield store, f"❌ {e}", "0 chunks", render_panel(steps), steps


def clear_all():
    return [], [], render_panel([]), ""

# ── UI ─────────────────────────────────────────────────────────────────────────

CSS = """
@import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;500&family=IBM+Plex+Sans:wght@300;400;500;600&display=swap');

body, .gradio-container {
  font-family: 'IBM Plex Sans', sans-serif !important;
  background: #f0f2f8 !important;
}
.gradio-container { max-width: 1300px !important; }

/* Topbar */
.app-topbar {
  background: #fff;
  border-bottom: 1px solid #e2e5f0;
  padding: 14px 20px;
  display: flex;
  align-items: center;
  gap: 12px;
  margin-bottom: 0;
}

/* Section labels */
.section-lbl {
  font-family: 'IBM Plex Mono', monospace;
  font-size: 10px;
  text-transform: uppercase;
  letter-spacing: 0.1em;
  color: #94a3b8;
  margin: 10px 0 4px;
}

/* Status */
.status-box textarea {
  font-family: 'IBM Plex Mono', monospace !important;
  font-size: 11.5px !important;
  line-height: 1.5 !important;
  background: #f8f9fc !important;
  border-color: #e2e5f0 !important;
}

/* Chatbot */
.chatbot-panel {
  border: 1px solid #e2e5f0 !important;
  border-radius: 10px !important;
  background: #fff !important;
  overflow: hidden;
}

/* Query input */
.query-wrap textarea {
  font-family: 'IBM Plex Mono', monospace !important;
  font-size: 13px !important;
  border-color: #e2e5f0 !important;
  background: #fff !important;
}
.query-wrap textarea:focus {
  border-color: #0891b2 !important;
}
"""

TOPBAR = """
<div class="app-topbar">
  <div style="display:flex;flex-direction:column;gap:2px">
    <div style="font-family:'IBM Plex Mono',monospace;font-size:16px;font-weight:600;color:#0f172a;letter-spacing:0.02em">
      rag_agent
    </div>
  </div>
</div>
"""

with gr.Blocks(title="RAG Agent") as demo:

    store_state = gr.State(EmbeddingStore())
    steps_state = gr.State([])           # list of step dicts driving the viz panel

    gr.HTML(TOPBAR)

    with gr.Row(equal_height=False):

        # ── Left sidebar ────────────────────────────────────────────────────
        with gr.Column(scale=1, min_width=260):
            gr.HTML('<div class="section-lbl">Document</div>')
            pdf_upload  = gr.File(label="Upload PDF", file_types=[".pdf"], type="filepath", show_label=False)
            process_btn = gr.Button("⚡  Embed PDF", variant="primary", size="sm")

            chunk_badge = gr.Textbox(
                value="0 chunks", interactive=False,
                lines=1, show_label=False,
                elem_classes="status-box",
            )
            pdf_status = gr.Textbox(
                value="Upload a PDF and click Embed.",
                interactive=False, lines=3,
                show_label=False, elem_classes="status-box",
            )

        # ── Center: chat ─────────────────────────────────────────────────────
        with gr.Column(scale=2):
            gr.HTML('<div class="section-lbl" style="margin-top:0">Chat</div>')
            chatbot = gr.Chatbot(
                height=440,
                show_label=False,
                placeholder=(
                    "**How to use**\n\n"
                    "① Upload a PDF → click **Embed PDF**\n"
                    "② Ask a question and watch the agent work →"
                ),
                render_markdown=True,
                elem_classes="chatbot-panel",
            )
            with gr.Row():
                query_box = gr.Textbox(
                    placeholder="Ask a question about your document…",
                    lines=2, scale=5, show_label=False,
                    elem_classes="query-wrap",
                )
                with gr.Column(scale=1, min_width=100):
                    send_btn  = gr.Button("Ask →", variant="primary")
                    clear_btn = gr.Button("Clear",  variant="secondary", size="sm")

        # ── Right: live agent viz ─────────────────────────────────────────────
        with gr.Column(scale=2):
            gr.HTML('<div class="section-lbl" style="margin-top:0">Agent decisions</div>')
            viz_panel = gr.HTML(value=EMPTY_PANEL)

    # ── Event wiring ──────────────────────────────────────────────────────────

    process_btn.click(
        fn=process_pdf,
        inputs=[pdf_upload, store_state, steps_state],
        outputs=[store_state, pdf_status, chunk_badge, viz_panel, steps_state],
    )

    def ask(query, store, history, steps):
        for h, s, panel in run_agent(query, store, history, steps):
            yield h, s, panel, ""

    send_btn.click(
        fn=ask,
        inputs=[query_box, store_state, chatbot, steps_state],
        outputs=[chatbot, steps_state, viz_panel, query_box],
        show_progress="hidden",
    )
    query_box.submit(
        fn=ask,
        inputs=[query_box, store_state, chatbot, steps_state],
        outputs=[chatbot, steps_state, viz_panel, query_box],
        show_progress="hidden",
    )
    clear_btn.click(
        fn=clear_all,
        outputs=[chatbot, steps_state, viz_panel, query_box],
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, css=CSS)