""" 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 # ---------------------------------------------------------------------- # Source title formatting (filename -> human-readable document title) # ---------------------------------------------------------------------- def source_title(filename): """Human-readable title from a source filename. Only removes file extensions and export-timestamp suffixes and tidies separators — it never invents words, so it cannot hallucinate a title. If the cleaned result is a single token with no spaces (an accession/ID code such as 's41467-022-31430-0' or 'nihms385546'), it is prefixed with 'Document ' so the model refers to a document rather than a bare code. Always returns a non-empty string. """ if not filename: return "Reference document" original = os.path.basename(str(filename)) name = re.sub(r"\.(pdf|png|jpe?g|txt|md)$", "", original, flags=re.IGNORECASE) # strip trailing export/tracking timestamp suffix, e.g. "-07-13-2026_04_20_PM" name = re.sub(r"[-_]\d{1,2}[-_]\d{1,2}[-_]\d{2,4}[_-]\d{1,2}[_-]\d{2}[_-][AP]M$", "", name) # tidy separators name = name.replace(" _ ", " — ") # site/section separator used by several files name = name.replace("+", " ").replace("_", " ") # hyphens joining two letters -> spaces (leave digit-joining hyphens in IDs alone) name = re.sub(r"(?<=[A-Za-z])-(?=[A-Za-z])", " ", name) name = re.sub(r"\s{2,}", " ", name).strip(" -—") if not name: return original if " " not in name: # single token => identifier/code return f"Document {name}" return name # ---------------------------------------------------------------------- # Participant-facing source links # ---------------------------------------------------------------------- # Keyed by the exact filename used in the RAG corpus. # If a source has a public/participant-accessible URL, add it here. # If it does not, leave it out or set it to None; its citation will remain # plain text rather than clickable. SOURCE_REGISTRY = { "Autologous Stem Cell Transplant for Multiple Myeloma (IMF, 2026).pdf": "https://www.myeloma.org/autologous-stem-cell-transplant", "Bispecific Therapies for Multiple Myeloma (International Myeloma Foundation, 2026).pdf": "https://www.myeloma.org/emerging-therapies/bispecific-therapies", "Blood Protein Tests for Multiple Myeloma (IMF, 2025).pdf": "https://www.myeloma.org/blood-protein-testing", "Bone Marrow Tests for Multiple Myeloma (IMF, 2025).pdf": "https://www.myeloma.org/bone-marrow-tests", "CAR T-Cell Therapy for Multiple Myeloma (International Myeloma Foundation, 2024).pdf": "https://www.myeloma.org/emerging-therapies/car-t-cell-therapy", "Comprehensive Molecular Profiling of Multiple Myeloma by African and European Ancestry (Manojlovic et al., 2017).pdf": "https://journals.plos.org/plosgenetics/article?id=10.1371/journal.pgen.1007087", "Durie-Salmon Staging for Multiple Myeloma (IMF, 2026).pdf": "https://www.myeloma.org/durie-salmon-staging", "FDA Approval of Teclistamab Plus Daratumumab for Relapsed Multiple Myeloma (2026).pdf": "https://www.fda.gov/drugs/resources-information-approved-drugs/fda-approves-teclistamab-combination-daratumumab-hyaluronidase-fihj-relapsed-or-refractory-multiple", "First-Line Treatment for Multiple Myeloma (IMF, 2026).pdf": "https://www.myeloma.org/frontline-treatment-options", "Genetic Heterogeneity and Drug Resistance in Relapsed Multiple Myeloma (Vo et al., 2022).pdf": "https://www.nature.com/articles/s41467-022-31430-0", "Genetic Heterogeneity in Multiple Myeloma and Targeted Therapy (Lohr et al., 2014).pdf": "https://www.cell.com/cancer-cell/fulltext/S1535-6108(13)00542-4?_returnURL=https%3A%2F%2Flinkinghub.elsevier.com%2Fretrieve%2Fpii%2FS1535610813005424%3Fshowall%3Dtrue", "Immunocytokine Therapy for Multiple Myeloma (International Myeloma Foundation, 2026).pdf": "https://www.myeloma.org/emerging-therapies/immunocytokine-therapy", "Initial Genome Sequencing of Multiple Myeloma (Chapman et al., 2011).pdf": "https://www.nature.com/articles/nature09837", "ISS and R-ISS Staging for Multiple Myeloma (IMF, 2026).pdf": "https://www.myeloma.org/international-staging-system-iss-reivised-iss-r-iss", "Kidney Function Tests for Multiple Myeloma (IMF, 2025).pdf": "https://www.myeloma.org/myeloma-kidney-tests", "M-Protein Tests for Multiple Myeloma (IMF, 2026).pdf": "https://www.myeloma.org/monoclonal-protein-tests", "Maintenance Therapy for Multiple Myeloma (IMF, 2024).pdf": "https://www.myeloma.org/maintenance-therapy", "Mayo Clinic mSMART Relapsed Myeloma Treatment Guidelines, Version 9 (2026).pdf": "https://static1.squarespace.com/static/5b44f08ac258b493a25098a3/t/69b45653e3e1365b6857a106/1773426259531/03_13_26_v9+Relapsed+MM+mSMART%2BRelapsedMM.pdf", "MGUS, Smoldering Myeloma, and Active Myeloma (IMF, 2026).pdf": "https://www.myeloma.org/what-are-mgus-smm-mm", "Minimal Residual Disease in Multiple Myeloma (IMF).pdf": "https://www.myeloma.org/videos/minimal-residual-disease-mrd", "Multiple Myeloma Autologous Stem Cell Transplantation Patient Toolkit (MMRF, 2026).pdf": "https://themmrf.org/wp-content/uploads/2026/03/Autologous-Stem-Cell-Transplantation-patient_toolkit-2026.pdf", "Multiple Myeloma Blood Tests (IMF, 2026).pdf": "https://www.myeloma.org/multiple-myeloma-blood-tests", "Multiple Myeloma Complications (IMF, 2025).pdf": "https://www.myeloma.org/medical-conditions", "Multiple Myeloma Diagnosis and Treatment (Mayo Clinic, 2024).pdf": "https://www.mayoclinic.org/diseases-conditions/multiple-myeloma/diagnosis-treatment/drc-20353383", "Multiple Myeloma FAQ (IMF, 2026).pdf": "https://www.myeloma.org/myeloma-cancer-questions", "Multiple Myeloma Imaging Studies (IMF, 2025).pdf": "https://www.myeloma.org/multiple-myeloma-imaging-studies", "Multiple Myeloma Immunotherapy Patient Toolkit (MMRF, 2026).pdf": "https://themmrf.org/wp-content/uploads/2026/03/Immunotherapy-patient_toolkit-2026.pdf", "Multiple Myeloma Overview (IMF, 2026).pdf": "https://www.myeloma.org/what-is-multiple-myeloma", "Multiple Myeloma Relapse (IMF, 2025).pdf": "https://www.myeloma.org/treatment/relapse-definition", "Multiple Myeloma Signs and Diagnosis (IMF, 2025).pdf": "https://www.myeloma.org/newly-diagnosed/do-you-have-myeloma", "Multiple Myeloma Staging and Risk Stratification (IMF, 2019).pdf": "https://www.myeloma.org/multiple-myeloma/multiple-myeloma-prognosis-staging", "Multiple Myeloma Testing (IMF, 2026).pdf": "https://www.myeloma.org/multiple-myeloma-tests", "Multiple Myeloma Testing, Staging and Prognosis (IMF, 2025).pdf": "https://www.myeloma.org/tests-staging", "NCCN Multiple Myeloma Patient Guidelines (2026).pdf": "https://www.nccn.org/patients/guidelines/content/PDF/myeloma-patient.pdf", "Perspectives on the Risk-Stratified Treatment of Multiple Myeloma (Davies et al., 2022).pdf": "https://pmc.ncbi.nlm.nih.gov/articles/PMC9894570/", "Quality of Life in Longitudinal Multiple Myeloma Studies (Nielsen et al., 2017).pdf": "https://onlinelibrary.wiley.com/doi/10.1111/ejh.12882", "Quality of Life in Multiple Myeloma Treatment (Sonneveld et al., 2013).pdf": "https://www.nature.com/articles/leu2013185", "Questions to Ask Your Myeloma Doctor (IMF).pdf": "https://www.myeloma.org/resource-library/tip-card-ask-your-doctor", "Routine Health Screenings for Myeloma Patients (IMF, 2021).pdf": "https://www.myeloma.org/health-screenings", "The Role of Minimal Residual Disease Testing in Myeloma Treatment Selection and Drug Development (Anderson et al., 2017).pdf": "https://aacrjournals.org/clincancerres/article/23/15/3980/257045/The-Role-of-Minimal-Residual-Disease-Testing-in", "Types of Multiple Myeloma (IMF, 2025).pdf": "https://www.myeloma.org/types-of-myeloma", } # ---------------------------------------------------------------------- # 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 def source_link_map(self): """ Map the human-readable document titles used in chatbot citations to participant-accessible URLs. Sources without a URL remain plain-text citations. """ links = {} for source in self.sources: url = SOURCE_REGISTRY.get(source) if url: links[source_title(source)] = url return links # ---- 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 source-attributed context block, labeling each passage with the document's readable title (not 'Source N').""" parts = [] for r in retrieved: title = source_title(r["source"]) parts.append(f"[Document: {title}]\n{r['text']}") return "\n\n---\n\n".join(parts)