Upload 4 files

app.py

@@ -1,36 +1,36 @@
 """
-P3 Myeloma - AI Follow-up Chatbot
-==================================
+P3 Myeloma - AI Follow-up Chatbot (knowledge-base grounded)
+===========================================================
 Takes a participant's answers to the 16 Shared Decision Making questions
 and runs an adaptive follow-up conversation using Google's Gemini API.
+
+GROUNDING:
+General informational questions from the user are answered STRICTLY from the
+documents in the uploaded ZIP corpus (Archive_2.zip) using retrieval-augmented
+generation. At startup the app extracts text from every PDF/PNG in the corpus,
+embeds it with Gemini embeddings, and builds a searchable index. Each user
+question retrieves the most relevant passages, and Gemini is instructed to
+answer only from those passages (or say the answer isn't in the documents).
+
 Deploy on Hugging Face Spaces (SDK: Gradio).
-Set your Gemini key as a Space Secret named GEMINI_API_KEY,
-or paste it in the UI field at runtime.
-
-NOTE ON HARDWARE:
-This app only calls the Gemini API and does NOT need a GPU. The recommended
-fix for "No @spaces.GPU function detected during startup" is to switch the
-Space hardware to **CPU basic** (Settings -> Hardware). If the Space must stay
-on ZeroGPU, the small `spaces.GPU` stub below satisfies the startup check.
+- Upload Archive_2.zip to the Space repo root (or set KB_ZIP / KB_DIR).
+- Set your Gemini key as a Space Secret named GEMINI_API_KEY.
+- Recommended packages.txt: poppler-utils, tesseract-ocr  (see notes at bottom).
 """
 
 import os
 import json
 
 # ----------------------------------------------------------------------
-# ZeroGPU compatibility shim.
-# On ZeroGPU hardware the runtime requires at least one @spaces.GPU function
-# to exist at startup, otherwise it raises:
-#   "No @spaces.GPU function detected during startup"
-# This app is API-only (no GPU work), so we register a trivial stub.
-# If `spaces` isn't installed (e.g. CPU basic), we skip it silently.
+# ZeroGPU compatibility shim (harmless on CPU basic; skipped if not present).
+# Satisfies "No @spaces.GPU function detected during startup" if the Space is
+# on ZeroGPU. This app is API-only and does not use a GPU — CPU basic is fine.
 # ----------------------------------------------------------------------
 try:
     import spaces
 
     @spaces.GPU
     def _zerogpu_warmup():
-        """No-op to satisfy the ZeroGPU startup check. Not used at runtime."""
         return None
 except Exception:
     pass
@@ -39,6 +39,8 @@ from google import genai
 from google.genai import types
 import gradio as gr
 
+from knowledge_base import KnowledgeBase
+
 # ----------------------------------------------------------------------
 # The 16 profile questions (verbatim from the intake form)
 # ----------------------------------------------------------------------
@@ -87,6 +89,27 @@ MAX_FOLLOWUPS = 8  # cap the adaptive conversation
 # The conversation always opens with this fixed question from the patient.
 FIRST_QUESTION = "Using my profile, what are the best treatment options for my myeloma?"
 
+# ----------------------------------------------------------------------
+# Knowledge base — built once at startup.
+# ----------------------------------------------------------------------
+KB = KnowledgeBase()
+_KB_BUILD_ATTEMPTED = False
+
+
+def ensure_kb(client):
+    """Build the KB on first successful client. Returns (ready, message)."""
+    global _KB_BUILD_ATTEMPTED
+    if KB.ready:
+        return True, ""
+    if _KB_BUILD_ATTEMPTED and KB.error:
+        return False, KB.error
+    _KB_BUILD_ATTEMPTED = True
+    KB.build(client)
+    if KB.ready:
+        return True, ""
+    return False, (KB.error or "Knowledge base could not be initialized.")
+
+
 # ----------------------------------------------------------------------
 # Prompt building
 # ----------------------------------------------------------------------
@@ -103,28 +126,20 @@ def detect_tensions(answers):
     a = {i: (answers[i].strip().lower() == "yes") for i in range(len(answers))}
     flags = []
 
-    # Aggressive cure-seeking vs quality-over-quantity
     if (a.get(5) or a.get(6) or a.get(7)) and a.get(11):
         flags.append("Wants an aggressive/curative approach but also values quality of life over quantity. Worth clarifying how they weigh these when they conflict.")
-    # Willing to endure side effects vs wants active lifestyle / quality of life
     if a.get(7) and (a.get(14) or a.get(11)):
         flags.append("Willing to endure many side effects, yet wants to stay active / prioritizes quality of life. Probe acceptable side-effect threshold.")
-    # Lives alone AND no strong social support -> logistics risk
     if a.get(3) and not a.get(4):
         flags.append("Lives alone with no caregiver and limited social support. Treatment logistics and safety monitoring need exploration.")
-    # Lives alone but prefers home meds -> safety of self-administration
     if a.get(3) and a.get(9):
         flags.append("Lives alone but prefers taking medications at home. Explore support for safe self-administration and side-effect monitoring.")
-    # Defers fully to team but also wants to collaborate / understand options
     if a.get(2) and (a.get(0) or a.get(1)):
         flags.append("Says they leave decisions to the care team, yet also wants to collaborate / understand options. Clarify how involved they actually want to be.")
-    # Cost matters + trial interest (trials can offset cost) — opportunity, not conflict
     if a.get(13) and a.get(12):
         flags.append("Cost matters and they're open to clinical trials — trials may reduce drug cost; worth surfacing.")
-    # Fewer pills preference but aggressive
     if a.get(10) and (a.get(6) or a.get(7)):
         flags.append("Prefers the fewest pills possible but wants an aggressive approach. Explore tolerance for treatment intensity vs convenience.")
-    # Outpatient + lives alone
     if a.get(8) and a.get(3):
         flags.append("Prefers outpatient treatment but lives alone — explore transport and post-visit support.")
     return flags
@@ -135,33 +150,49 @@ def system_instruction():
         "You are a warm, plain-spoken health navigator helping a multiple myeloma patient "
         "(relapsed or refractory) prepare for a shared decision-making conversation with their "
         "care team. You are NOT a doctor and you never give medical advice, diagnoses, dosing, or "
-        "treatment recommendations. Your job is ONLY to ask thoughtful follow-up questions that help "
-        "the patient clarify their own values, priorities, constraints, and concerns.\n\n"
-        "RULES:\n"
-        "1. Ask exactly ONE question per turn. Keep it to 1-3 sentences, conversational, jargon-free.\n"
-        "2. Base each question on the patient's profile and their previous answers. Prioritize the "
+        "treatment recommendations. Your job is to (a) ask thoughtful follow-up questions that help "
+        "the patient clarify their own values, priorities, constraints, and concerns, and (b) answer "
+        "the patient's general informational questions using ONLY the reference documents provided to "
+        "you in the RETRIEVED CONTEXT for that turn.\n\n"
+        "STRICT GROUNDING RULES (most important):\n"
+        "A. For any general/informational question the patient asks (e.g. 'what is CAR T-cell "
+        "therapy?', 'what are the side effects of stem cell transplant?', 'what does relapsed mean?'), "
+        "you MUST base your answer solely on the text inside the RETRIEVED CONTEXT block for that turn. "
+        "Do NOT use outside knowledge, and do NOT add facts that are not present in that context.\n"
+        "B. If the RETRIEVED CONTEXT does not contain enough information to answer, say so plainly: "
+        "\"I couldn't find that in the reference documents provided for this tool. Please ask your care "
+        "team.\" Do not guess or fill gaps from general knowledge.\n"
+        "C. When you use information from the context, attribute it naturally to the source document "
+        "name(s) shown in the context (e.g. \"According to the NCCN patient guideline ...\").\n"
+        "D. Never invent drug names, doses, statistics, or study results that are not in the context.\n\n"
+        "CONVERSATION RULES:\n"
+        "1. When you ask a follow-up question, ask exactly ONE question per turn, 1-3 sentences, "
+        "conversational and jargon-free.\n"
+        "2. Base follow-up questions on the patient's profile and previous answers. Prioritize the "
         "FLAGGED TENSIONS provided — gently explore apparent contradictions without judgment.\n"
         "3. Do not repeat questions already asked. Build on what they say.\n"
-        "4. Never recommend or rank treatments. If the patient asks for medical advice, kindly "
-        "redirect them to discuss it with their care team, then continue with a follow-up question.\n"
+        "4. Never recommend or rank treatments. If the patient asks for medical advice or 'what should "
+        "I do', kindly redirect them to their care team, then (if appropriate) continue with a follow-up "
+        "question.\n"
         "5. Tone: empathetic, respectful, never alarming.\n"
         "6. TREATMENT-OPTIONS QUESTION: When the patient asks what the best treatment options are for "
         "their myeloma (e.g. 'Using my profile, what are the best treatment options for my myeloma?'), "
-        "your answer must be grounded STRICTLY in the 16 yes/no profile answers provided to you. Do "
+        "combine TWO sources: their 16 yes/no profile answers AND the RETRIEVED CONTEXT documents. Do "
         "this:\n"
         "   - Open with one warm sentence.\n"
-        "   - Then walk through how THEIR specific stated priorities map to the kinds of treatment "
-        "approaches commonly discussed for relapsed/refractory myeloma. Tie each point explicitly back "
-        "to their answers (for example: a preference for taking medications at home and the fewest pills "
-        "points toward asking the team about convenient outpatient or oral regimens; openness to clinical "
-        "trials means trials are worth raising; willingness to endure side effects and wanting an "
-        "aggressive approach points toward asking about more intensive options; prioritizing quality of "
-        "life points toward asking about regimens that protect daily functioning; living alone with "
-        "limited support points toward asking about logistics and monitoring).\n"
-        "   - Do NOT invent clinical details, drug names beyond broad well-known categories, dosing, or "
-        "anything not implied by their answers. Do NOT rank, prescribe, or state which option is "
-        "medically best. Frame everything as 'based on what you told us, here are options and questions "
-        "to raise with your care team.'\n"
+        "   - Walk through how THEIR specific stated priorities map to the kinds of treatment approaches "
+        "described in the RETRIEVED CONTEXT for relapsed/refractory myeloma. Tie each point explicitly "
+        "back to their answers (for example: a preference for medications at home and fewest pills points "
+        "toward asking about convenient outpatient or oral regimens; openness to clinical trials means "
+        "trials are worth raising; willingness to endure side effects and wanting an aggressive approach "
+        "points toward asking about more intensive options; prioritizing quality of life points toward "
+        "regimens that protect daily functioning; living alone with limited support points toward "
+        "logistics and monitoring).\n"
+        "   - Only describe treatment approaches that actually appear in the RETRIEVED CONTEXT. Do NOT "
+        "invent clinical details, drug names, dosing, or anything not supported by the context or implied "
+        "by their answers. Do NOT rank, prescribe, or state which option is medically best. Frame "
+        "everything as 'based on what you told us and these reference materials, here are options and "
+        "questions to raise with your care team.'\n"
         "   - End by reminding them their care team must confirm what is medically appropriate, then ask "
         "your first single follow-up question drawn from their profile and the flagged tensions.\n"
         "7. When you judge that you have enough to summarize their priorities (or after several "
@@ -177,7 +208,7 @@ def system_instruction():
 def get_client(user_key):
     key = (user_key or "").strip() or os.environ.get("GEMINI_API_KEY", "").strip()
     if not key:
-        return None, "No API key found. Paste your Gemini API key above, or set GEMINI_API_KEY as a Space secret."
+        return None, "No API key found. Set GEMINI_API_KEY as a Space secret."
     try:
         client = genai.Client(api_key=key)
         return client, None
@@ -195,16 +226,54 @@ def to_gemini_contents(history):
     return contents
 
 
+def _latest_user_query(convo):
+    for m in reversed(convo):
+        if m["role"] == "user":
+            return m["content"]
+    return ""
+
+
 def gemini_reply(client, convo, n_followups):
-    """convo: internal history list. Returns assistant text."""
+    """convo: internal history list. Retrieves KB context for the latest user
+    turn and injects it, then asks Gemini to answer grounded in that context."""
+    # Retrieve relevant passages for the most recent user message.
+    context_block = ""
+    query = _latest_user_query(convo)
+    if KB.ready and query.strip():
+        try:
+            retrieved = KB.retrieve(client, query)
+            if retrieved:
+                context_block = KB.context_block(retrieved)
+        except Exception:
+            context_block = ""
+
     contents = to_gemini_contents(convo)
-    # Nudge toward summarizing as we approach the cap
-    extra = ""
+
+    # Build the per-turn system instruction with the retrieved context appended.
+    sys = system_instruction()
+    if context_block:
+        sys += (
+            "\n\n==================== RETRIEVED CONTEXT (reference documents) ====================\n"
+            "The following passages were retrieved from the uploaded reference documents for THIS "
+            "turn. Answer general/informational questions using ONLY these passages. If they do not "
+            "contain the answer, say you couldn't find it in the reference documents.\n\n"
+            f"{context_block}\n"
+            "================================================================================\n"
+        )
+    else:
+        sys += (
+            "\n\n[No reference passages were retrieved for this turn. If the patient asked a general "
+            "informational question, tell them you couldn't find it in the reference documents and "
+            "suggest they ask their care team. You may still ask a values-clarifying follow-up "
+            "question or work with their profile.]\n"
+        )
+
     if n_followups >= MAX_FOLLOWUPS - 1:
-        extra = "\n\nYou have asked enough questions. Provide the final <<SUMMARY>> now."
+        sys += "\n\nYou have asked enough questions. Provide the final <<SUMMARY>> now."
+
     cfg = types.GenerateContentConfig(
-        system_instruction=system_instruction() + extra,
-        temperature=0.7,
+        system_instruction=sys,
+        temperature=0.3,          # lower temp -> stays closer to the sources
         max_output_tokens=2048,
     )
     resp = client.models.generate_content(model=MODEL, contents=contents, config=cfg)
@@ -226,12 +295,16 @@ def start_chat(api_key, *radio_values):
         gr.Warning(err)
         return (gr.update(), [], [], 0, gr.update(visible=True), gr.update(visible=False))
 
+    # Build the knowledge base on first use (may take a bit on cold start).
+    ready, msg = ensure_kb(client)
+    if not ready:
+        gr.Warning(f"Reference documents unavailable: {msg}")
+        # We still allow the conversation, but answers will note missing docs.
+
     profile = build_profile_text(answers)
     flags = detect_tensions(answers)
     flag_text = "\n".join(f"- {f}" for f in flags) if flags else "- (No obvious contradictions; explore their highest-stakes priorities.)"
 
-    # Seed the conversation with the profile as the first user turn so every later
-    # turn (including when the user pastes the sample question) can be grounded in it.
     seed = (
         "Here is the patient's completed profile (16 yes/no answers):\n\n"
         f"{profile}\n\n"
@@ -251,12 +324,12 @@ def start_chat(api_key, *radio_values):
     convo.append({"role": "assistant", "content": reply})
     chat_display = [{"role": "assistant", "content": reply}]
     return (
-        chat_display,          # chatbot
-        chat_display,          # display_state
-        convo,                 # convo_state
-        0,                     # followup count (no follow-up asked yet)
-        gr.update(visible=False),  # hide intake
-        gr.update(visible=True),   # show chat
+        chat_display,
+        chat_display,
+        convo,
+        0,
+        gr.update(visible=False),
+        gr.update(visible=True),
     )
 
 
@@ -270,6 +343,8 @@ def respond(user_msg, chat_display, convo, n_followups, api_key):
         gr.Warning(err)
         return chat_display, chat_display, convo, n_followups, user_msg
 
+    ensure_kb(client)  # no-op if already built
+
     convo = convo + [{"role": "user", "content": user_msg}]
     chat_display = chat_display + [{"role": "user", "content": user_msg}]
 
@@ -296,19 +371,18 @@ def respond(user_msg, chat_display, convo, n_followups, api_key):
 def reset_all():
     radio_resets = [gr.update(value=None) for _ in QUESTIONS]
     return (
-        [],   # chatbot
-        [],   # display_state
-        [],   # convo_state
-        0,    # followups
-        gr.update(visible=True),   # intake
-        gr.update(visible=False),  # chat
+        [],
+        [],
+        [],
+        0,
+        gr.update(visible=True),
+        gr.update(visible=False),
         *radio_resets,
     )
 
 
 CSS = """
 @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
-/* ---- Base: white background, black text, professional font ---- */
 .gradio-container, .gradio-container * {
     font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Helvetica, Arial, sans-serif !important;
 }
@@ -322,7 +396,6 @@ body, .gradio-container .prose, .gradio-container p,
 .gradio-container span, .gradio-container label {
     color: #111111 !important;
 }
-/* ---- Header ---- */
 #app-header {
     border-bottom: 1px solid #e6e6e6;
     padding: 4px 0 18px 0;
@@ -341,7 +414,6 @@ body, .gradio-container .prose, .gradio-container p,
     color: #5a5a5a;
     margin: 4px 0 0 0;
 }
-/* ---- Question cards ---- */
 .q-card {
     border: 1px solid #e2e2e2 !important;
     border-radius: 8px !important;
@@ -354,12 +426,10 @@ body, .gradio-container .prose, .gradio-container p,
     color: #111111 !important;
     font-weight: 500 !important;
 }
-/* ---- Radio choices ---- */
 .gradio-container input[type="radio"] + span,
 .gradio-container .wrap label {
     color: #111111 !important;
 }
-/* ---- Buttons ---- */
 button.primary, .gradio-container button.primary {
     background: #7a0c2e !important;
     color: #ffffff !important;
@@ -377,7 +447,6 @@ button.secondary, .gradio-container button.secondary {
     border-radius: 6px !important;
     font-weight: 600 !important;
 }
-/* ---- Chatbot ---- */
 .gradio-container .chatbot, .gradio-container [class*="chatbot"] {
     background: #ffffff !important;
     border: 1px solid #e2e2e2 !important;
@@ -386,7 +455,6 @@ button.secondary, .gradio-container button.secondary {
 .gradio-container .message.bot, .gradio-container .message.user {
     color: #111111 !important;
 }
-/* ---- Inputs ---- */
 .gradio-container textarea, .gradio-container input[type="text"],
 .gradio-container input[type="password"] {
     background: #ffffff !important;
@@ -394,11 +462,9 @@ button.secondary, .gradio-container button.secondary {
     border: 1px solid #cfcfcf !important;
     border-radius: 6px !important;
 }
-/* ---- Accordion ---- */
 .gradio-container .label-wrap, .gradio-container details summary {
     color: #111111 !important;
 }
-/* ---- Sample question hint ---- */
 .sample-q {
     background: #f7f7f7 !important;
     border: 1px solid #e2e2e2 !important;
@@ -439,14 +505,11 @@ with gr.Blocks(title="P3 Myeloma Follow-up Chatbot", css=CSS,
             "<p id='app-subtitle'>Follow-up Questions for Shared Decision Making</p>"
         )
 
-    # API key is read from the Space secret GEMINI_API_KEY (hidden, not user-facing).
     api_key = gr.State("")
-
     convo_state = gr.State([])
     display_state = gr.State([])
     followups = gr.State(0)
 
-    # ---- Intake panel ----
     with gr.Group(visible=True) as intake_panel:
         gr.Markdown("**Answer all 16 questions, then click Start Follow-up.**")
         radios = []
@@ -456,7 +519,6 @@ with gr.Blocks(title="P3 Myeloma Follow-up Chatbot", css=CSS,
                 radios.append(r)
         start_btn = gr.Button("Start Follow-up", variant="primary")
 
-    # ---- Chat panel ----
     with gr.Group(visible=False) as chat_panel:
         chatbot = gr.Chatbot(label="Follow-up Conversation", type="messages", height=460)
         gr.HTML(
@@ -475,7 +537,6 @@ with gr.Blocks(title="P3 Myeloma Follow-up Chatbot", css=CSS,
         inputs=[api_key] + radios,
         outputs=[chatbot, display_state, convo_state, followups, intake_panel, chat_panel],
     )
-
     send_btn.click(
         respond,
         inputs=[msg, display_state, convo_state, followups, api_key],
@@ -486,7 +547,6 @@ with gr.Blocks(title="P3 Myeloma Follow-up Chatbot", css=CSS,
         inputs=[msg, display_state, convo_state, followups, api_key],
         outputs=[chatbot, display_state, convo_state, followups, msg],
     )
-
     restart_btn.click(
         reset_all,
         inputs=None,
@@ -494,4 +554,19 @@ with gr.Blocks(title="P3 Myeloma Follow-up Chatbot", css=CSS,
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()
+
+# ----------------------------------------------------------------------
+# HUGGING FACE SPACE SETUP NOTES
+# ----------------------------------------------------------------------
+# 1) Files in the Space repo:
+#      app.py
+#      knowledge_base.py
+#      Archive_2.zip          <- the uploaded corpus (repo root)
+#      requirements.txt       <- google-genai, gradio, numpy, pypdf, pdf2image, pytesseract, pillow
+#      packages.txt           <- poppler-utils, tesseract-ocr   (system deps for PDF text + OCR)
+# 2) Space Secret:
+#      GEMINI_API_KEY = <your key>
+# 3) Hardware: CPU basic is sufficient (no GPU needed).
+# 4) On first launch the app extracts + embeds the corpus once and caches the
+#    index to _kb_index.npz, so subsequent restarts are fast.

knowledge_base.py

@@ -0,0 +1,320 @@
+"""
+knowledge_base.py — Retrieval-Augmented Generation over the uploaded ZIP corpus.
+
+At import/startup:
+  1. Locate the corpus (an extracted folder, or a .zip we extract once).
+  2. Extract text from every PDF (poppler pdftotext -> pypdf fallback) and
+     OCR any PNG/JPG (optional; skipped gracefully if tesseract is absent).
+  3. Chunk the text, embed all chunks once with Gemini embeddings, cache to disk.
+
+At query time:
+  - embed the user question, cosine-retrieve top-K chunks, and return them so
+    the chat model can answer STRICTLY from that context.
+"""
+import os, re, glob, json, hashlib, subprocess, zipfile, shutil, time
+import numpy as np
+from google.genai import types
+
+EMBED_MODEL = "gemini-embedding-001"
+EMBED_DIM = 768                    # output_dimensionality for the embedding model
+TOP_K = 8                          # chunks retrieved per query
+TARGET_CHARS = 1400
+OVERLAP = 200
+MACOSX = "__MACOSX"
+
+# Where to look for the corpus, in priority order. On HF Spaces the repo root
+# is the app's working dir, so an uploaded "Archive_2.zip" or a pre-extracted
+# "knowledge_base/" folder both work.
+def _corpus_candidates():
+    return [os.environ.get("KB_DIR", "").strip(), "knowledge_base", "kb", "data", "."]
+
+def _zip_candidates():
+    return [os.environ.get("KB_ZIP", "").strip(), "Archive_2.zip"]
+EXTRACT_DIR = "_kb_extracted"
+CACHE_PATH = "_kb_index.npz"
+MANIFEST_PATH = "_kb_manifest.json"
+
+
+# ----------------------------------------------------------------------
+# Corpus location / extraction
+# ----------------------------------------------------------------------
+def _has_docs(d):
+    if not d or not os.path.isdir(d):
+        return False
+    for p in glob.glob(os.path.join(d, "**", "*"), recursive=True):
+        if MACOSX in p or os.path.basename(p).startswith("._"):
+            continue
+        if p.lower().endswith((".pdf", ".png", ".jpg", ".jpeg", ".txt", ".md")):
+            return True
+    return False
+
+
+def locate_corpus():
+    """Return a directory containing the source documents, extracting a zip if needed."""
+    for d in _corpus_candidates():
+        if _has_docs(d):
+            # Avoid treating '.' as corpus if it only contains app files;
+            # require at least one PDF for the '.' case.
+            if d == "." and not glob.glob("*.pdf"):
+                continue
+            return d
+    for z in _zip_candidates():
+        if z and os.path.isfile(z):
+            os.makedirs(EXTRACT_DIR, exist_ok=True)
+            with zipfile.ZipFile(z) as zf:
+                for name in zf.namelist():
+                    if MACOSX in name or os.path.basename(name).startswith("._"):
+                        continue
+                    zf.extract(name, EXTRACT_DIR)
+            if _has_docs(EXTRACT_DIR):
+                return EXTRACT_DIR
+    return None
+
+
+def list_source_files(kb_dir):
+    out = []
+    for p in sorted(glob.glob(os.path.join(kb_dir, "**", "*"), recursive=True)):
+        if MACOSX in p or os.path.basename(p).startswith("._"):
+            continue
+        if os.path.isfile(p) and p.lower().endswith((".pdf", ".png", ".jpg", ".jpeg", ".txt", ".md")):
+            out.append(p)
+    return out
+
+
+# ----------------------------------------------------------------------
+# Text extraction
+# ----------------------------------------------------------------------
+def _pdftotext(path):
+    try:
+        r = subprocess.run(["pdftotext", "-layout", path, "-"],
+                           capture_output=True, timeout=120)
+        if r.returncode == 0:
+            return r.stdout.decode("utf-8", "ignore")
+    except Exception:
+        pass
+    return ""
+
+
+def _pypdf_text(path):
+    try:
+        from pypdf import PdfReader
+        reader = PdfReader(path)
+        return "\n".join((pg.extract_text() or "") for pg in reader.pages)
+    except Exception:
+        return ""
+
+
+def _ocr_image(path):
+    try:
+        import pytesseract
+        from PIL import Image
+        return pytesseract.image_to_string(Image.open(path).convert("RGB"))
+    except Exception:
+        return ""
+
+
+def _ocr_pdf(path):
+    """Rasterize + OCR a scanned/image PDF. Optional; empty if tools missing."""
+    try:
+        import pytesseract
+        from pdf2image import convert_from_path
+        text = []
+        for img in convert_from_path(path, dpi=200):
+            text.append(pytesseract.image_to_string(img))
+        return "\n".join(text)
+    except Exception:
+        return ""
+
+
+def extract_file_text(path):
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".pdf":
+        txt = _pdftotext(path)
+        if len(txt.strip()) < 50:
+            txt = _pypdf_text(path)
+        if len(txt.strip()) < 50:
+            txt = _ocr_pdf(path)          # image-only PDF
+        return txt
+    if ext in (".png", ".jpg", ".jpeg"):
+        return _ocr_image(path)
+    if ext in (".txt", ".md"):
+        try:
+            with open(path, encoding="utf-8", errors="ignore") as fh:
+                return fh.read()
+        except Exception:
+            return ""
+    return ""
+
+
+# ----------------------------------------------------------------------
+# Chunking
+# ----------------------------------------------------------------------
+def clean_text(t):
+    t = t.replace("\x00", " ")
+    t = re.sub(r"[ \t]+", " ", t)
+    t = re.sub(r"\n{3,}", "\n\n", t)
+    return t.strip()
+
+
+def chunk_text(text, source, target_chars=TARGET_CHARS, overlap=OVERLAP):
+    text = clean_text(text)
+    if not text:
+        return []
+    chunks, start, n = [], 0, len(text)
+    while start < n:
+        end = min(start + target_chars, n)
+        if end < n:
+            window = text[start:end]
+            cut = max(window.rfind("\n\n"), window.rfind(". "), window.rfind("\n"))
+            if cut > target_chars * 0.5:
+                end = start + cut + 1
+        chunk = text[start:end].strip()
+        if len(chunk) > 40:
+            chunks.append({"text": chunk, "source": os.path.basename(source)})
+        if end >= n:
+            break
+        start = max(end - overlap, start + 1)
+    return chunks
+
+
+# ----------------------------------------------------------------------
+# Embedding (Gemini) with batching + disk cache
+# ----------------------------------------------------------------------
+def _l2norm(v):
+    v = np.asarray(v, dtype=np.float32)
+    nrm = np.linalg.norm(v, axis=-1, keepdims=True)
+    nrm[nrm == 0] = 1.0
+    return v / nrm
+
+
+def _embed_batch(client, texts, task_type):
+    """Return np.array (len(texts), EMBED_DIM). Falls back to per-item on error."""
+    cfg = types.EmbedContentConfig(task_type=task_type, output_dimensionality=EMBED_DIM)
+    try:
+        resp = client.models.embed_content(model=EMBED_MODEL, contents=texts, config=cfg)
+        return np.array([e.values for e in resp.embeddings], dtype=np.float32)
+    except Exception:
+        vecs = []
+        for t in texts:
+            resp = client.models.embed_content(model=EMBED_MODEL, contents=t, config=cfg)
+            vecs.append(resp.embeddings[0].values)
+        return np.array(vecs, dtype=np.float32)
+
+
+def _corpus_fingerprint(files):
+    h = hashlib.sha256()
+    for p in sorted(files):
+        st = os.stat(p)
+        h.update(p.encode()); h.update(str(st.st_size).encode()); h.update(str(int(st.st_mtime)).encode())
+    h.update(f"{EMBED_MODEL}:{EMBED_DIM}:{TARGET_CHARS}:{OVERLAP}".encode())
+    return h.hexdigest()
+
+
+class KnowledgeBase:
+    def __init__(self):
+        self.ready = False
+        self.error = None
+        self.chunks = []          # list of {text, source}
+        self.matrix = None        # (N, EMBED_DIM) normalized
+        self.sources = []         # unique source filenames
+        self.dir = None
+
+    # ---- build or load ----
+    def build(self, client, embed_batch=_embed_batch, log=print):
+        try:
+            self.dir = locate_corpus()
+            if not self.dir:
+                self.error = ("No knowledge-base documents found. Upload Archive_2.zip to the "
+                              "Space (repo root) or add a knowledge_base/ folder of PDFs.")
+                return False
+            files = list_source_files(self.dir)
+            if not files:
+                self.error = "Knowledge-base folder found but contains no readable documents."
+                return False
+
+            fp = _corpus_fingerprint(files)
+            if self._load_cache(fp):
+                self.ready = True
+                log(f"[KB] Loaded cached index: {len(self.chunks)} chunks from {len(self.sources)} files.")
+                return True
+
+            # Extract + chunk
+            all_chunks = []
+            for p in files:
+                txt = extract_file_text(p)
+                cs = chunk_text(txt, p)
+                if cs:
+                    all_chunks.extend(cs)
+                    log(f"[KB]  {os.path.basename(p)}: {len(cs)} chunks")
+                else:
+                    log(f"[KB]  {os.path.basename(p)}: no extractable text (skipped)")
+            if not all_chunks:
+                self.error = "Documents found but no text could be extracted from any of them."
+                return False
+
+            # Embed in batches
+            texts = [c["text"] for c in all_chunks]
+            vecs = []
+            B = 64
+            for i in range(0, len(texts), B):
+                batch = texts[i:i + B]
+                vecs.append(embed_batch(client, batch, "RETRIEVAL_DOCUMENT"))
+                log(f"[KB]  embedded {min(i+B, len(texts))}/{len(texts)} chunks")
+            matrix = np.vstack(vecs)
+            self.chunks = all_chunks
+            self.matrix = _l2norm(matrix)
+            self.sources = sorted({c["source"] for c in all_chunks})
+            self._save_cache(fp)
+            self.ready = True
+            log(f"[KB] Index built: {len(self.chunks)} chunks from {len(self.sources)} files.")
+            return True
+        except Exception as e:
+            self.error = f"Failed to build knowledge base: {e}"
+            return False
+
+    def _save_cache(self, fp):
+        try:
+            np.savez_compressed(CACHE_PATH, matrix=self.matrix,
+                                texts=np.array([c["text"] for c in self.chunks], dtype=object),
+                                srcs=np.array([c["source"] for c in self.chunks], dtype=object),
+                                fp=np.array([fp]))
+            with open(MANIFEST_PATH, "w") as fh:
+                json.dump({"fingerprint": fp, "sources": self.sources,
+                           "n_chunks": len(self.chunks)}, fh, indent=2)
+        except Exception:
+            pass
+
+    def _load_cache(self, fp):
+        if not os.path.isfile(CACHE_PATH):
+            return False
+        try:
+            d = np.load(CACHE_PATH, allow_pickle=True)
+            if str(d["fp"][0]) != fp:
+                return False
+            self.matrix = d["matrix"].astype(np.float32)
+            texts = list(d["texts"]); srcs = list(d["srcs"])
+            self.chunks = [{"text": t, "source": s} for t, s in zip(texts, srcs)]
+            self.sources = sorted(set(srcs))
+            return True
+        except Exception:
+            return False
+
+    # ---- retrieve ----
+    def retrieve(self, client, query, k=TOP_K, embed_batch=_embed_batch):
+        if not self.ready:
+            return []
+        qv = embed_batch(client, [query], "RETRIEVAL_QUERY")[0]
+        qv = _l2norm(qv)
+        sims = self.matrix @ qv
+        k = min(k, len(sims))
+        idx = np.argpartition(-sims, k - 1)[:k]
+        idx = idx[np.argsort(-sims[idx])]
+        return [{"text": self.chunks[i]["text"], "source": self.chunks[i]["source"],
+                 "score": float(sims[i])} for i in idx]
+
+    def context_block(self, retrieved):
+        """Format retrieved chunks as a numbered, source-attributed context block."""
+        parts = []
+        for i, r in enumerate(retrieved, 1):
+            parts.append(f"[Source {i}: {r['source']}]\n{r['text']}")
+        return "\n\n---\n\n".join(parts)

packages.txt

@@ -0,0 +1,2 @@
+poppler-utils
+tesseract-ocr

requirements.txt

@@ -1,3 +1,7 @@
-gradio==5.49.1
-google-genai==1.38.0
-pydantic==2.11.10
+gradio==5.9.1
+google-genai>=0.3.0
+numpy>=1.24
+pypdf>=4.0
+pdf2image>=1.17
+pytesseract>=0.3.10
+pillow>=10.0