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 = """
"""
EMPTY_PANEL = PANEL_CSS + """
Agent decisions will appear here in real-time
Upload a PDF → ask a question
"""
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'── {esc(s.get("label",""))} ──
')
continue
elif t == "think":
body = f'{esc(s["text"])}
'
elif t == "search":
body = (
f'Querying corpus for:
'
f'{esc(s["query"])}
'
+ (f'k = {s["k"]}
' 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"""
"""
body = (
f''
f'Found {len(results)} relevant chunk{"s" if len(results)!=1 else ""}
'
f'{chunk_html}
'
)
elif t == "fetch":
body = (
f'Fetching full text of:
'
f'{esc(s["chunk_id"])}
'
)
if s.get("text"):
preview = s["text"][:300] + ("…" if len(s["text"]) > 300 else "")
body += f'{esc(preview)}
'
elif t == "list":
body = f'Scanning all {s.get("total",0)} chunks in corpus for context…
'
elif t == "answer":
body = (
f'Answer synthesized
'
f'from {s.get("chunks_used",0)} retrieved chunk(s) across {s.get("steps",0)} retrieval step(s)
'
)
elif t == "error":
body = f'{esc(s["text"])}
'
elif t == "embed":
body = f'{esc(s["text"])}
'
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"""
{body}
""")
return PANEL_CSS + '' + "".join(cards) + "
"
# ── 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 = """
"""
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('Document
')
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('Chat
')
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('Agent decisions
')
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)