knowledge_base.py

from __future__ import annotations

import os
import io
import torch
import fitz  # PyMuPDF
from PIL import Image

from chunking import normalize_extracted_text, split_text_into_chunks
from evidence import stable_id

from retrieval_utils import select_scored_results

############################
# Load & Configure Retrieval
############################

def _pick_dtype(device: str):
    if device.startswith("cuda"):
        return torch.float16
    # bf16 is nice if CPU supports it, but safest is float32 on CPU
    return torch.float32


def load_retrieval_model(model_choice="colpali", device="cpu"):
    """
    Backward-compatible loader with extra, faster VLM options.
    Returns: (model, processor, model_type)
    model_type equals model_choice, preserving your external behavior.

    Supported choices (additive over your original):
      - "siglip"  : google/siglip-base-patch16-224 (fast, multimodal)
      - "clip"    : openai/clip-vit-base-patch32 (fast, multimodal)
      - "colpali" : vidore/colpali-v1.2-hf (your original heavy retriever)
      - "all-minilm" : all-MiniLM-L6-v2 (your original text-only)
    """
    dtype = _pick_dtype(device)

    if model_choice == "siglip":
        from transformers import SiglipModel, SiglipProcessor
        name = "google/siglip-base-patch16-224"
        model = SiglipModel.from_pretrained(name, torch_dtype=dtype).to(device).eval()
        processor = SiglipProcessor.from_pretrained(name)
        model_type = "siglip"

    elif model_choice == "clip":
        from transformers import CLIPModel, CLIPProcessor
        name = "openai/clip-vit-base-patch32"
        model = CLIPModel.from_pretrained(name, torch_dtype=dtype).to(device).eval()
        processor = CLIPProcessor.from_pretrained(name)
        model_type = "clip"

    elif model_choice == "colpali":
        from transformers import ColPaliForRetrieval, ColPaliProcessor
        name = "vidore/colpali-v1.2-hf"
        # keep your semantics, just ensure dtype/device
        model = ColPaliForRetrieval.from_pretrained(name, torch_dtype=torch.bfloat16).to(device).eval()
        processor = ColPaliProcessor.from_pretrained(name)
        model_type = "colpali"

    elif model_choice == "all-minilm":
        from sentence_transformers import SentenceTransformer
        model = SentenceTransformer("all-MiniLM-L6-v2", device=device)
        processor = None
        model_type = "all-minilm"

    else:
        raise ValueError(f"Unsupported retrieval model choice: {model_choice}")

    return model, processor, model_type


def _l2norm(x: torch.Tensor) -> torch.Tensor:
    x = x.float()
    return x / (x.norm(dim=-1, keepdim=True) + 1e-12)


def _as_flat_float_tensor(value, *, label: str) -> torch.Tensor:
    """Normalize embedding-like values into a non-empty flat float tensor."""
    if value is None:
        raise ValueError(f"{label} embedding is missing")
    if not isinstance(value, torch.Tensor):
        value = torch.as_tensor(value)
    flat = value.detach().float().cpu().view(-1)
    if flat.numel() == 0:
        raise ValueError(f"{label} embedding is empty")
    return flat


def validate_embedding_compatibility(query_emb, doc_emb) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Ensure query and document embeddings can be scored in the same vector space.

    The audit found that hybrid retrieval paths can silently mix dimensions. Make
    that state impossible to score: callers either get compatible tensors or a
    concrete error explaining the incompatible shapes.
    """
    q_vec = _as_flat_float_tensor(query_emb, label="query")
    d_vec = _as_flat_float_tensor(doc_emb, label="document")
    if q_vec.numel() != d_vec.numel():
        raise ValueError(
            "Embedding dimension mismatch: "
            f"query has {q_vec.numel()} values, document has {d_vec.numel()} values"
        )
    return q_vec, d_vec


def embed_text(query, model, processor, model_type="colpali", device="cpu"):
    """
    Backward-compatible. Always returns a torch.Tensor on CPU for uniformity.
    """
    with torch.no_grad():
        if model_type == "colpali":
            # Keep your original pathway (text → embeddings)
            inputs = processor(text=[query], truncation=True, max_length=512, return_tensors="pt").to(device)
            outputs = model(**inputs)
            emb = outputs.embeddings.mean(dim=1).squeeze(0)
            return _l2norm(emb).cpu()

        elif model_type == "all-minilm":
            emb = model.encode(query, convert_to_tensor=True)
            return _l2norm(emb).cpu()

        elif model_type == "siglip":
            # SigLIP provides aligned text/image spaces
            # get_text_features is available via forward helpers in HF >= 4.40
            inputs = processor(text=[query], return_tensors="pt", padding=True, truncation=True, max_length=64).to(device)
            try:
                emb = model.get_text_features(**inputs)
            except AttributeError:
                # fallback to forward and pool
                out = model(**inputs)
                emb = out.text_embeds
            emb = emb.squeeze(0)
            return _l2norm(emb).cpu()

        elif model_type == "clip":
            inputs = processor(text=[query], return_tensors="pt", padding=True, truncation=True, max_length=77).to(device)
            try:
                emb = model.get_text_features(**inputs)
            except AttributeError:
                out = model(**inputs)
                emb = out.text_embeds
            emb = emb.squeeze(0)
            return _l2norm(emb).cpu()

        else:
            raise ValueError(f"Unsupported model_type: {model_type}")


##################
# Scoring & Search
##################

def late_interaction_score(query_emb, doc_emb):
    q_vec, d_vec = validate_embedding_compatibility(query_emb, doc_emb)
    q_norm = q_vec / (q_vec.norm() + 1e-12)
    d_norm = d_vec / (d_vec.norm() + 1e-12)
    return float(torch.dot(q_norm, d_norm))


def retrieve(
    query,
    corpus,
    model,
    processor,
    top_k=3,
    model_type="colpali",
    device="cpu",
    text_model=None,
    max_chunks_per_doc_in_results=3,
):
    if top_k <= 0 or not corpus:
        return []

    # Keep query and document embeddings in the active retriever's vector space.
    # The optional text_model argument remains for API compatibility but is no
    # longer used to avoid mixed SigLIP/CLIP vs MiniLM dimensions.
    query_embedding = embed_text(query, model, processor, model_type=model_type, device=device)
    scores = []
    for idx, entry in enumerate(corpus):
        try:
            score = late_interaction_score(query_embedding, entry['embedding'])
        except ValueError as exc:
            metadata = entry.get("metadata", {})
            source = metadata.get("file_path") or metadata.get("url") or f"corpus[{idx}]"
            raise ValueError(f"Incompatible embedding for {source}: {exc}") from exc
        scores.append(score)
    return select_scored_results(
        corpus,
        scores,
        top_k=top_k,
        max_per_document=max_chunks_per_doc_in_results,
    )


##################################
# Building a Corpus from a Folder
##################################

def _pool_mean(tensors):
    if not tensors:
        return None
    stacked = torch.stack([t.float() for t in tensors], dim=0)
    return _l2norm(stacked.mean(dim=0))


def _embed_long_text(text: str, model, processor, model_type: str, device: str, max_len=1200, stride=800):
    """Chunk long text to keep memory bounded; mean-pool chunk embeddings."""
    # Adjust max_len based on model type
    if model_type in ["siglip", "clip"]:
        max_len = 200  # Much shorter for vision models
        stride = 150
    elif model_type == "colpali":
        max_len = 400  # Medium for ColPali
        stride = 300
    
    chunks = []
    i = 0
    while i < len(text):
        chunk = text[i:i+max_len]
        i += stride
        chunks.append(chunk)
    embs = []
    for c in chunks:
        embs.append(embed_text(c, model, processor, model_type=model_type, device=device))
    return _pool_mean(embs)


def _embed_image(img: Image.Image, model, processor, model_type: str, device: str):
    with torch.no_grad():
        if model_type == "siglip":
            inputs = processor(images=img, return_tensors="pt").to(device)
            try:
                feats = model.get_image_features(**inputs)
            except AttributeError:
                out = model(**inputs)
                feats = out.image_embeds
            return _l2norm(feats.squeeze(0)).cpu()
        elif model_type == "clip":
            inputs = processor(images=img, return_tensors="pt").to(device)
            try:
                feats = model.get_image_features(**inputs)
            except AttributeError:
                out = model(**inputs)
                feats = out.image_embeds
            return _l2norm(feats.squeeze(0)).cpu()
        else:
            return None  # non-vision models


def _pdf_pages_to_images(pdf_path: str, max_pages: int = 4, dpi: int = 144):
    try:
        doc = fitz.open(pdf_path)
        pages = []
        for i, page in enumerate(doc):
            if i >= max_pages:
                break
            pix = page.get_pixmap(dpi=dpi)
            img = Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB")
            pages.append(img)
        return pages
    except Exception:
        return []


# Text formats the README documents as supported. Keep this list and the README in
# sync: an unlisted extension is skipped silently, which looks like empty retrieval.
PLAIN_TEXT_EXTENSIONS = (".txt", ".md", ".markdown", ".py", ".json", ".yaml", ".yml", ".csv", ".rst")
IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg")


def load_corpus_from_dir(
    corpus_dir,
    model,
    processor,
    device="cpu",
    model_type="colpali",
    chunk_chars=1200,
    chunk_overlap=200,
    max_chunks_per_doc=40,
):
    """
    Scan 'corpus_dir' for txt, pdf, and image files, embed their content,
    and return a list of { 'embedding': torch.Tensor(cpu), 'metadata':... }.
    - For VLMs (siglip/clip): images & PDF pages are embedded directly (no OCR needed).
    - For text-only models: text is extracted (OCR for images as fallback) and embedded.
    - PDFs: combine (mean-pool) text chunks + first few rendered pages (VLMs) for robustness.
    """
    corpus = []
    if not corpus_dir or not os.path.isdir(corpus_dir):
        return corpus
    skipped = []

    for filename in os.listdir(corpus_dir):
        file_path = os.path.join(corpus_dir, filename)
        if not os.path.isfile(file_path):
            continue

        ext = filename.lower()
        text = ""
        img_embs = []

        # --- Plain text formats ---
        if ext.endswith(PLAIN_TEXT_EXTENSIONS):
            try:
                with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
                    text = f.read()
            except Exception as e:
                print(f"[WARN] Failed to read text file {file_path}: {e}")
                continue

        # --- PDF ---
        elif ext.endswith(".pdf"):
            # Extract text quickly
            try:
                doc = fitz.open(file_path)
                tparts = []
                for i, page in enumerate(doc):
                    if i >= 10:  # cap for speed
                        break
                    t = page.get_text("text").strip()
                    if not t:
                        t = page.get_text("blocks").strip() or ""
                    if t:
                        tparts.append(t)
                text = "\n".join(tparts)
            except Exception as e:
                print(f"[WARN] Failed to read PDF {file_path}: {e}")
                text = ""

            # For VLMs, also embed first few pages as images (fast, no OCR)
            if model_type in ("siglip", "clip"):
                for img in _pdf_pages_to_images(file_path, max_pages=4):
                    e = _embed_image(img, model, processor, model_type, device)
                    if e is not None:
                        img_embs.append(e)

        # --- Images ---
        elif ext.endswith(IMAGE_EXTENSIONS):
            if model_type in ("siglip", "clip"):
                try:
                    img = Image.open(file_path).convert("RGB")
                    e = _embed_image(img, model, processor, model_type, device)
                    if e is not None:
                        img_embs.append(e)
                except Exception as e:
                    print(f"[WARN] Image load failed {file_path}: {e}")
                    continue
            else:
                # Text-only models: use OCR fallback
                try:
                    import pytesseract
                    img = Image.open(file_path)
                    text = pytesseract.image_to_string(img)
                except Exception as e:
                    print(f"[WARN] OCR failed for image {file_path}: {e}")
                    continue
        else:
            # Skipping silently here reads downstream as "retrieval found nothing",
            # which is very hard to diagnose from the report. Say so out loud.
            skipped.append(filename)
            continue

        # Build entries. Text is indexed as chunks so retrieved local evidence carries
        # real body text into the report; a single pooled vector could only ever supply
        # a 100-character snippet, which reads downstream as "no evidence".
        try:
            entries_before = len(corpus)
            canonical = normalize_extracted_text(text) if text.strip() else ""
            if canonical:
                evidence_id = stable_id(file_path, prefix="local")
                chunks = split_text_into_chunks(
                    evidence_id,
                    canonical,
                    chunk_chars=chunk_chars,
                    chunk_overlap=chunk_overlap,
                    max_chunks_per_doc=max_chunks_per_doc,
                )
                for chunk in chunks:
                    chunk_emb = _embed_long_text(chunk.text, model, processor, model_type, device)
                    if chunk_emb is None:
                        continue
                    corpus.append({
                        "embedding": chunk_emb.cpu(),
                        "metadata": {
                            "file_path": file_path,
                            "type": "localchunk",
                            "evidence_id": f"{evidence_id}#chunk-{chunk.chunk_index}",
                            "parent_evidence_id": evidence_id,
                            "chunk_index": chunk.chunk_index,
                            "char_start": chunk.char_start,
                            "char_end": chunk.char_end,
                            "text": chunk.text,
                            "snippet": chunk.text[:240],
                        },
                    })

            if img_embs:
                corpus.append({
                    "embedding": _pool_mean(img_embs).cpu(),
                    "metadata": {
                        "file_path": file_path,
                        "type": "local",
                        "snippet": (text[:100].replace('\n', ' ') + "...") if text else "",
                    },
                })

            if len(corpus) == entries_before:
                # Nothing usable
                continue
        except Exception as e:
            print(f"[WARN] Skipping embedding for local file {file_path} due to error: {e}")

    if skipped:
        print(
            f"[WARN] Skipped {len(skipped)} unsupported file(s) in {corpus_dir}: "
            f"{', '.join(sorted(skipped)[:8])}"
            f"{' ...' if len(skipped) > 8 else ''}"
        )
    print(f"[INFO] Loaded {len(corpus)} document(s) from {corpus_dir}")
    return corpus


###########################
# KnowledgeBase Class (API)
###########################

class KnowledgeBase:
    """
    Same public API with explicit embedding-compatibility checks.
    """
    def __init__(self, model, processor, model_type="colpali", device="cpu", text_model=None,
                 max_chunks_per_doc_in_results=3):
        self.model = model
        self.processor = processor
        self.model_type = model_type
        self.device = device
        self.text_model = text_model  # Kept for backward-compatible constructor calls.
        self.corpus = []
        self.max_chunks_per_doc_in_results = max_chunks_per_doc_in_results

    def add_documents(self, entries):
        self.corpus.extend(entries)

    def search(self, query, top_k=3, max_chunks_per_doc_in_results=None):
        return retrieve(
            query,
            self.corpus,
            self.model,
            self.processor,
            top_k=top_k,
            model_type=self.model_type,
            device=self.device,
            text_model=self.text_model,
            max_chunks_per_doc_in_results=(
                self.max_chunks_per_doc_in_results
                if max_chunks_per_doc_in_results is None
                else max_chunks_per_doc_in_results
            ),
        )