extractor.py

"""Local PDF extraction: text, images, and query-aware snippets."""

import os
import re
from pathlib import Path
from typing import Dict, List, Optional

import fitz  # PyMuPDF

from config import CFG
from retrieval import build_snippet, retrieve_relevant_passages


def select_relevant_snippet(
    text: str,
    query: str,
    window: Optional[int] = None,
) -> str:
    window = window or CFG.max_chars_per_paper
    if not text:
        return ""
    if not query or len(text) <= window:
        return text[:window]

    passages = retrieve_relevant_passages(text, query, top_k=CFG.retrieval_top_k)
    snippet = build_snippet(passages, max_chars=window)
    return snippet or text[:window]


def extract_pdf(
    pdf_path: str,
    query: str = "",
    *,
    analyze_images: Optional[bool] = None,
    image_query: str = "",
) -> Dict:
    """Extract text, selected images, and a query-aware snippet from a PDF."""
    doc = fitz.open(pdf_path)
    text_parts: List[str] = []
    images: List[Dict] = []
    img_dir = os.path.join(CFG.cache_dir, "images", Path(pdf_path).stem)
    if CFG.extract_images:
        Path(img_dir).mkdir(parents=True, exist_ok=True)

    img_count = 0
    try:
        for page_index, page in enumerate(doc):
            text_parts.append(page.get_text("text") or "")
            if not CFG.extract_images or img_count >= CFG.max_images_per_paper:
                continue
            for image_index, image in enumerate(page.get_images(full=True)):
                if img_count >= CFG.max_images_per_paper:
                    break
                try:
                    pix = fitz.Pixmap(doc, image[0])
                    if pix.width < 80 or pix.height < 80:
                        pix = None
                        continue
                    if pix.n - pix.alpha > 3:
                        pix = fitz.Pixmap(fitz.csRGB, pix)
                    out_path = os.path.join(
                        img_dir,
                        f"p{page_index}_{image_index}.png",
                    )
                    pix.save(out_path)
                    images.append(
                        {
                            "page": page_index,
                            "path": out_path,
                            "w": pix.width,
                            "h": pix.height,
                        }
                    )
                    img_count += 1
                    pix = None
                except Exception as exc:  # noqa: BLE001 - skip one bad image
                    print(f"[extractor] img skip p{page_index} #{image_index}: {exc}")
    finally:
        doc.close()

    full_text = "\n".join(text_parts)
    retrieval_passages = (
        retrieve_relevant_passages(full_text, query, top_k=CFG.retrieval_top_k)
        if query
        else []
    )
    snippet = (
        build_snippet(retrieval_passages, max_chars=CFG.max_chars_per_paper)
        if retrieval_passages
        else select_relevant_snippet(full_text, query)
    )
    do_analyze_images = CFG.analyze_images_default if analyze_images is None else analyze_images
    image_analyses = []
    if do_analyze_images and images:
        from image_understanding import analyze_images as _analyze_images

        image_analyses = _analyze_images(images, prompt=image_query or query)

    return {
        "ok": True,
        "pdf_path": pdf_path,
        "full_text": full_text,
        "full_text_len": len(full_text),
        "snippet": snippet,
        "retrieval_method": CFG.retrieval_method,
        "retrieval_passages": retrieval_passages,
        "images": images,
        "image_analyses": image_analyses,
    }