Jonathan KirkwoodandClaude Fable 5
1dde915540
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
"""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
|
|
.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"}
|
|
|
|
|
|
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 _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 == ".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_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."""
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
manifest: list[dict] = []
|
|
if not os.path.isdir(inbox):
|
|
return manifest
|
|
seen: dict[str, int] = {}
|
|
for fn in sorted(os.listdir(inbox)):
|
|
src = os.path.join(inbox, fn)
|
|
if not os.path.isfile(src):
|
|
continue
|
|
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
|