""" 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 # ---------------------------------------------------------------------- # 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 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)