knowledge_base.py

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