boardroom-map/orchestrator/extraction.py

"""Document text extraction — runs on the StartOS box (CPU only).

Confidential documents are dropped into /data/inbox. Before anything is shipped
to the Sparks, we extract plain text here so that only normalized text (never the
original binaries) crosses to the review containers. Supported formats:

    .pdf            -> pypdf
    .pptx           -> python-pptx (text frames, tables, chart data, notes)
    .docx           -> python-docx
    .txt .md .text  -> read as UTF-8

Anything else is skipped with a note. Each extracted document becomes a single
UTF-8 .txt file in the per-job staging directory, which is rsynced to the Spark
and mounted read-only into every reviewer container at /docs.
"""
from __future__ import annotations

import os

TEXT_EXTS = {".txt", ".md", ".text", ".markdown"}
SUPPORTED = TEXT_EXTS | {".pdf", ".docx", ".pptx"}


def _extract_pdf(path: str) -> str:
    from pypdf import PdfReader

    reader = PdfReader(path)
    parts = []
    for i, page in enumerate(reader.pages, 1):
        try:
            txt = page.extract_text() or ""
        except Exception as e:
            txt = f"(page {i}: extraction error: {e})"
        parts.append(f"\n\n===== page {i} =====\n{txt}")
    return "".join(parts).strip()


def _extract_docx(path: str) -> str:
    import docx

    doc = docx.Document(path)
    lines = [p.text for p in doc.paragraphs]
    # Include table cell text too — contracts/specs often hide content in tables.
    for table in doc.tables:
        for row in table.rows:
            cells = [c.text.strip() for c in row.cells]
            if any(cells):
                lines.append(" | ".join(cells))
    return "\n".join(lines).strip()


def _pptx_chart_text(shape) -> list[str]:
    """Best-effort chart title + series names/values (chart XML varies wildly)."""
    lines: list[str] = []
    try:
        chart = shape.chart
        try:
            if chart.has_title and chart.chart_title.has_text_frame:
                lines.append(f"[chart] {chart.chart_title.text_frame.text}")
        except Exception:
            pass
        for plot in chart.plots:
            try:
                cats = [str(c) for c in (plot.categories or [])]
            except Exception:
                cats = []
            for series in plot.series:
                try:
                    name = str(series.name)
                except Exception:
                    name = "(series)"
                try:
                    vals = ["" if v is None else f"{v:g}" for v in series.values]
                except Exception:
                    vals = []
                if cats and len(cats) == len(vals):
                    pairs = ", ".join(f"{c}={v}" for c, v in zip(cats, vals))
                else:
                    pairs = ", ".join(vals)
                lines.append(f"[chart series] {name}: {pairs}")
    except Exception:
        pass
    return lines


def _extract_pptx(path: str) -> str:
    from pptx import Presentation

    prs = Presentation(path)
    parts: list[str] = []
    for n, slide in enumerate(prs.slides, 1):
        body: list[str] = []
        for shape in slide.shapes:
            if getattr(shape, "has_text_frame", False):
                txt = shape.text_frame.text.strip()
                if txt:
                    body.append(txt)
            if getattr(shape, "has_table", False):
                for row in shape.table.rows:
                    cells = [c.text.strip() for c in row.cells]
                    if any(cells):
                        body.append(" | ".join(cells))
            if getattr(shape, "has_chart", False):
                body.extend(_pptx_chart_text(shape))
        notes = ""
        try:
            if slide.has_notes_slide:
                notes = (slide.notes_slide.notes_text_frame.text or "").strip()
        except Exception:
            pass
        if notes:
            body.append(f"--- notes ---\n{notes}")
        if len("".join(body)) < 20:
            body.append(f"[low_text_slide: slide {n}]")
        parts.append(f"\n\n===== slide {n} =====\n" + "\n".join(body))
    return "".join(parts).strip()


def _extract_text(path: str) -> str:
    with open(path, errors="replace") as f:
        return f.read().strip()


def extract_file(path: str) -> str:
    ext = os.path.splitext(path)[1].lower()
    if ext == ".pdf":
        return _extract_pdf(path)
    if ext == ".pptx":
        return _extract_pptx(path)
    if ext == ".docx":
        return _extract_docx(path)
    if ext in TEXT_EXTS:
        return _extract_text(path)
    raise ValueError(f"unsupported file type: {ext or '(none)'}")


def _safe_name(name: str) -> str:
    base = os.path.splitext(os.path.basename(name))[0]
    keep = "".join(c if (c.isalnum() or c in "-_ ") else "_" for c in base).strip()
    return (keep or "document").replace(" ", "_")


def extract_files(files: list[str], out_dir: str, log=print) -> list[dict]:
    """Extract an explicit list of files (a deck unit) to .txt files in `out_dir`.

    Returns a manifest: [{source, out, chars, ok, error}]. Skips unsupported
    files (recorded with ok=False) rather than failing the whole job."""
    os.makedirs(out_dir, exist_ok=True)
    manifest: list[dict] = []
    seen: dict[str, int] = {}
    for src in files:
        fn = os.path.basename(src)
        ext = os.path.splitext(fn)[1].lower()
        rec = {"source": fn, "out": None, "chars": 0, "ok": False, "error": ""}
        if ext not in SUPPORTED:
            rec["error"] = f"unsupported type {ext or '(none)'}"
            log(f"[extract] skip {fn}: {rec['error']}")
            manifest.append(rec)
            continue
        try:
            text = extract_file(src)
        except Exception as e:
            rec["error"] = str(e)[:300]
            log(f"[extract] FAILED {fn}: {rec['error']}")
            manifest.append(rec)
            continue
        stem = _safe_name(fn)
        if stem in seen:
            seen[stem] += 1
            stem = f"{stem}-{seen[stem]}"
        else:
            seen[stem] = 1
        out_name = f"{stem}.txt"
        out_path = os.path.join(out_dir, out_name)
        header = f"# Source document: {fn}\n\n"
        with open(out_path, "w") as f:
            f.write(header + text + "\n")
        rec.update({"out": out_name, "chars": len(text), "ok": True})
        log(f"[extract] {fn} -> {out_name} ({len(text)} chars)")
        manifest.append(rec)
    return manifest


def extract_inbox(inbox: str, out_dir: str, log=print) -> list[dict]:
    """Extract every supported file in `inbox` to a .txt in `out_dir`.

    Returns a manifest: [{source, out, chars, ok, error}]. Skips unsupported
    files (recorded with ok=False) rather than failing the whole job."""
    if not os.path.isdir(inbox):
        os.makedirs(out_dir, exist_ok=True)
        return []
    files = [os.path.join(inbox, fn) for fn in sorted(os.listdir(inbox))
             if os.path.isfile(os.path.join(inbox, fn))]
    return extract_files(files, out_dir, log)