Create app.py

app.py

@@ -0,0 +1,484 @@
+import gradio as gr
+import fitz           # pymupdf
+import math
+import re
+import json
+import time
+from google import genai
+from google.genai import types
+
+# ─── Chunking ─────────────────────────────────────────────────────────────────
+
+CHUNK_SIZE    = 800    # chars — bigger is fine, embedding-2 handles 8192 tokens
+CHUNK_OVERLAP = 120
+
+
+def split_into_chunks(text: str) -> list:
+    paragraphs = [p.strip() for p in re.split(r"\n{2,}", text) if p.strip()]
+    chunks = []
+    for para in paragraphs:
+        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_similarity(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:
+    """
+    In-memory dense vector store backed by gemini-embedding-2-preview.
+    Uses 1536-dim MRL truncation: excellent quality, half the storage of 3072.
+    """
+    EMBED_MODEL = "gemini-embedding-2-preview"
+    EMBED_DIM   = 1536
+    BATCH_SIZE  = 20
+
+    def __init__(self):
+        self.chunks   = []   # list of {id, text, embedding}
+        self.client   = None
+        self._api_key = ""
+
+    def set_client(self, api_key):
+        if api_key != self._api_key:
+            self.client   = genai.Client(api_key=api_key)
+            self._api_key = api_key
+
+    def clear(self):
+        self.chunks = []
+
+    def _embed_batch(self, texts):
+        response = self.client.models.embed_content(
+            model=self.EMBED_MODEL,
+            contents=texts,
+            config=types.EmbedContentConfig(
+                task_type="RETRIEVAL_DOCUMENT",
+                output_dimensionality=self.EMBED_DIM,
+            ),
+        )
+        return [list(e.values) for e in response.embeddings]
+
+    def _embed_query(self, query):
+        response = self.client.models.embed_content(
+            model=self.EMBED_MODEL,
+            contents=query,
+            config=types.EmbedContentConfig(
+                task_type="RETRIEVAL_QUERY",
+                output_dimensionality=self.EMBED_DIM,
+            ),
+        )
+        return list(response.embeddings[0].values)
+
+    def add_chunks(self, texts, progress_cb=None):
+        start_idx = len(self.chunks)
+        total     = len(texts)
+        added     = 0
+
+        for batch_start in range(0, total, self.BATCH_SIZE):
+            batch  = texts[batch_start : batch_start + self.BATCH_SIZE]
+            embeds = self._embed_batch(batch)
+
+            for text, emb in zip(batch, embeds):
+                cid = f"chunk_{start_idx + added}"
+                self.chunks.append({"id": cid, "text": text, "embedding": emb})
+                added += 1
+
+            if progress_cb:
+                progress_cb(added, total)
+
+            if batch_start + self.BATCH_SIZE < total:
+                time.sleep(0.3)   # rate-limit buffer
+
+        return added
+
+    def search(self, query, k=5):
+        if not self.chunks:
+            return []
+        qvec   = self._embed_query(query)
+        scored = [
+            {"id": c["id"], "text": c["text"],
+             "score": cosine_similarity(qvec, c["embedding"])}
+            for c in self.chunks
+        ]
+        scored.sort(key=lambda x: x["score"], reverse=True)
+        return scored[:k]
+
+    def get_by_id(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_text(pdf_path):
+    doc   = fitz.open(pdf_path)
+    pages = [page.get_text() for page in doc]
+    doc.close()
+    return "\n\n".join(pages)
+
+
+# ─── Agent Tool Declarations ──────────────────────────────────────────────────
+
+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 to find different aspects."
+        ),
+        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 the full text of a specific chunk by its 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 an overview of all chunks (IDs + first 120 chars each). Use to understand corpus scope.",
+        parameters=types.Schema(type=types.Type.OBJECT, properties={}),
+    ),
+]
+
+AGENT_TOOLS  = [types.Tool(function_declarations=TOOL_DECLARATIONS)]
+
+SYSTEM_PROMPT = """You are a precise RAG (Retrieval-Augmented Generation) agent.
+
+You have access to a PDF corpus via semantic search tools powered by Gemini Embedding 2 (dense vectors).
+
+To answer questions:
+1. PLAN what information you need
+2. RETRIEVE using search_chunks — run multiple searches with different phrasings for multi-part questions
+3. EXPAND with get_chunk_by_id when you need the full text of a promising chunk
+4. SYNTHESIZE a final answer grounded only in retrieved content
+
+Rules:
+- Always retrieve before answering — never rely on prior knowledge alone
+- Cite chunk IDs inline e.g. [chunk_2] for every factual claim
+- If retrieved chunks lack sufficient information, say so clearly
+- Be thorough but concise"""
+
+
+# ─── Agentic Loop ─────────────────────────────────────────────────────────────
+
+def run_agent(api_key, query, store, history, log):
+    if not api_key.strip():
+        yield history, log + "\n❌ Enter your Gemini API key.", store
+        return
+    if not store or store.size == 0:
+        yield history, log + "\n❌ No PDF loaded. Upload and embed one first.", store
+        return
+    if not query.strip():
+        return
+
+    store.set_client(api_key.strip())
+    client = genai.Client(api_key=api_key.strip())
+
+    history = history + [{"role": "user", "content": query}]
+    log_out = log + f"\n\n{'─'*52}\n🔍  {query}\n"
+    yield history, log_out, store
+
+    # Build Gemini message list
+    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
+
+    while step < MAX_STEPS:
+        step += 1
+        log_out += f"\n⚙️  step {step}\n"
+        yield history, log_out, store
+
+        response = client.models.generate_content(
+            model="gemini-2.0-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_call  = False
+        tool_responses = []
+        text_parts     = []
+
+        for part in parts:
+            if hasattr(part, "text") and part.text and part.text.strip():
+                log_out += f"\n💭  {part.text.strip()[:400]}\n"
+                yield history, log_out, store
+                text_parts.append(part.text)
+
+            if hasattr(part, "function_call") and part.function_call:
+                has_tool_call = True
+                fn   = part.function_call
+                name = fn.name
+                args = dict(fn.args) if fn.args else {}
+
+                log_out += f"\n🔧  {name}({json.dumps(args)})\n"
+                yield history, log_out, store
+
+                if name == "search_chunks":
+                    results = store.search(args.get("query", ""), k=int(args.get("k", 4)))
+                    result  = {"results": [
+                        {"id": r["id"], "score": round(r["score"], 4), "text": r["text"]}
+                        for r in results
+                    ]}
+                    log_out += f"   ↳ {len(results)} chunks: {', '.join(r['id'] for r in results)}\n"
+
+                elif name == "get_chunk_by_id":
+                    chunk  = store.get_by_id(args.get("chunk_id", ""))
+                    result = ({"id": chunk["id"], "text": chunk["text"]}
+                              if chunk else {"error": "Chunk not found"})
+                    log_out += f"   ↳ fetched {args.get('chunk_id')}\n"
+
+                elif name == "list_all_chunks":
+                    result  = {"chunks": store.overview()}
+                    log_out += f"   ↳ {store.size} chunks listed\n"
+
+                else:
+                    result = {"error": f"Unknown tool: {name}"}
+
+                yield history, log_out, store
+
+                tool_responses.append(types.Part(
+                    function_response=types.FunctionResponse(name=name, response=result)
+                ))
+
+        if has_tool_call and tool_responses:
+            gemini_msgs.append(types.Content(role="user", parts=tool_responses))
+            continue
+
+        # No tool calls → final answer
+        final = "\n".join(text_parts).strip() or "(No response generated)"
+        history = history + [{"role": "assistant", "content": final}]
+        log_out += "\n✅  Answer ready\n"
+        yield history, log_out, store
+        return
+
+    log_out += f"\n⚠️  Max steps ({MAX_STEPS}) reached.\n"
+    yield history, log_out, store
+
+
+# ─── PDF Processing (generator for live status) ───────────────────────────────
+
+def process_pdf(pdf_file, api_key, store):
+    if pdf_file is None:
+        yield store, "⚠️ No file uploaded.", "0 chunks", gr.update()
+        return
+    if not api_key.strip():
+        yield store, "❌ Enter your Gemini API key first.", "0 chunks", gr.update()
+        return
+
+    try:
+        new_store = EmbeddingStore()
+        new_store.set_client(api_key.strip())
+
+        yield new_store, "📄 Extracting text from PDF…", "…", gr.update(interactive=False)
+        text = extract_pdf_text(pdf_file.name)
+        if not text.strip():
+            yield new_store, "❌ No extractable text in this PDF.", "0 chunks", gr.update(interactive=True)
+            return
+
+        char_count = len(text)
+
+        yield new_store, f"✂️  Splitting into chunks…", "…", gr.update(interactive=False)
+        chunks = split_into_chunks(text)
+        total  = len(chunks)
+
+        yield new_store, f"🔢 Embedding {total} chunks via gemini-embedding-2-preview…\n(this may take a moment)", f"0 / {total}", gr.update(interactive=False)
+
+        added = new_store.add_chunks(chunks)
+
+        msg = (
+            f"✅ Ready!  {char_count:,} chars  →  {added} chunks\n"
+            f"   Model : gemini-embedding-2-preview\n"
+            f"   Dims  : {new_store.EMBED_DIM} (MRL truncated from 3072)"
+        )
+        yield new_store, msg, f"{added} chunks", gr.update(interactive=True)
+
+    except Exception as e:
+        yield store, f"❌ {e}", "0 chunks", gr.update(interactive=True)
+
+
+def clear_chat():
+    return [], ""
+
+
+# ─── UI ───────────────────────────────────────────────────────────────────────
+
+CSS = """
+@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&display=swap');
+body, .gradio-container { font-family: 'IBM Plex Sans', sans-serif !important; }
+.status-out textarea, .chunk-badge textarea {
+    font-family: 'IBM Plex Mono', monospace !important;
+    font-size: 12px !important;
+}
+.log-out textarea {
+    font-family: 'IBM Plex Mono', monospace !important;
+    font-size: 11px !important;
+    color: #666 !important;
+}
+"""
+
+HEADER = """
+<div style="padding:16px 4px 12px; border-bottom:1px solid #e5e7eb; margin-bottom:8px">
+  <h1 style="font-family:'IBM Plex Mono',monospace; font-size:18px; font-weight:500; color:#111; margin:0">
+    rag_agent
+  </h1>
+  <p style="font-size:11px; color:#999; margin:4px 0 0; font-family:'IBM Plex Mono',monospace">
+    gemini-embedding-2-preview &nbsp;·&nbsp; 1536-dim dense retrieval &nbsp;·&nbsp; gemini-2.0-flash agentic loop
+  </p>
+</div>
+"""
+
+with gr.Blocks(title="RAG Agent") as demo:
+
+    store_state = gr.State(EmbeddingStore())
+
+    gr.HTML(HEADER)
+
+    with gr.Row(equal_height=False):
+
+        # Left panel
+        with gr.Column(scale=1, min_width=290):
+
+            api_key = gr.Textbox(
+                label="Gemini API Key",
+                placeholder="AIza…",
+                type="password",
+                lines=1,
+            )
+            pdf_upload = gr.File(
+                label="Upload PDF",
+                file_types=[".pdf"],
+                type="filepath",
+            )
+            process_btn = gr.Button("⚙  Embed PDF", variant="primary")
+            chunk_badge = gr.Textbox(
+                label="Index",
+                value="0 chunks",
+                interactive=False,
+                lines=1,
+                elem_classes="chunk-badge",
+            )
+            pdf_status = gr.Textbox(
+                label="Status",
+                value="Upload a PDF and click Embed.",
+                interactive=False,
+                lines=4,
+                elem_classes="status-out",
+            )
+            with gr.Accordion("Agent trace log", open=False):
+                agent_log = gr.Textbox(
+                    label="",
+                    value="",
+                    lines=22,
+                    max_lines=400,
+                    interactive=False,
+                    elem_classes="log-out",
+                )
+
+        # Right panel — chat
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(
+                label="",
+                height=520,
+                placeholder=(
+                    "**How to use**\n\n"
+                    "1. Paste your Gemini API key\n"
+                    "2. Upload a PDF → click **Embed PDF**\n"
+                    "   (chunks get embedded via `gemini-embedding-2-preview`)\n"
+                    "3. Ask questions — the agent retrieves, reasons, and answers with citations"
+                ),
+                render_markdown=True,
+                show_label=False,
+            )
+            with gr.Row():
+                query_box = gr.Textbox(
+                    placeholder="Ask a question about your document…",
+                    label="",
+                    lines=2,
+                    scale=5,
+                    show_label=False,
+                )
+                with gr.Column(scale=1, min_width=110):
+                    send_btn  = gr.Button("Ask →", variant="primary")
+                    clear_btn = gr.Button("Clear", variant="secondary")
+
+    # Wiring
+    process_btn.click(
+        fn=process_pdf,
+        inputs=[pdf_upload, api_key, store_state],
+        outputs=[store_state, pdf_status, chunk_badge, send_btn],
+    )
+
+    def ask(api_key, query, store, history, log):
+        yield from run_agent(api_key, query, store, history, log)
+
+    send_btn.click(
+        fn=ask,
+        inputs=[api_key, query_box, store_state, chatbot, agent_log],
+        outputs=[chatbot, agent_log, store_state],
+        show_progress="hidden",
+    )
+    query_box.submit(
+        fn=ask,
+        inputs=[api_key, query_box, store_state, chatbot, agent_log],
+        outputs=[chatbot, agent_log, store_state],
+        show_progress="hidden",
+    )
+    clear_btn.click(fn=clear_chat, outputs=[chatbot, agent_log])
+
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, css=CSS)