Document AI: Turn PDFs, Tables and Scans into Clean RAG Chunks

"""Turns the PDFs in pdfs/ into clean, self-describing chunks for a RAG index."""
import base64
import os
import re

import pypdf

import harness

HEADING = re.compile(r"^\d+\. [A-Z]")


# ---------- Step 1: read the pages ----------
def read_pages(folder="pdfs"):
    """One dict per page of every PDF in the folder, in file-name order: {"doc", "page" (from 1), "text"}."""
    # TODO (Step 1): for every .pdf in sorted(os.listdir(folder)), open it with pypdf.PdfReader and add one dict per
    # page: {"doc": file name, "page": number from 1, "text": page.extract_text() or ""}.
    raise NotImplementedError("Step 1: write read_pages()")


def needs_ocr(page):
    """True when the page has no usable text layer (fewer than 20 characters once stripped)."""
    # TODO (Step 1): True when the stripped text is shorter than 20 characters
    raise NotImplementedError("Step 1: write needs_ocr()")


# ---------- Step 2: clean the text ----------
def strip_running_lines(pages):
    """Copies of the pages without running headers and footers: in a document with two or more pages, a line that
    appears on every page (digits ignored, so "Page 1 of 2" and "Page 2 of 2" count as the same line) is removed."""
    raise NotImplementedError("strip_running_lines() arrives in Step 2")


def unwrap(text):
    """Undo the line wrapping: join words hyphenated at a line end ("comp-" + "leted" -> "completed"), then join each
    line to the one before with a space, except headings ("3. Pets") and the line after a heading, which start new lines."""
    raise NotImplementedError("unwrap() arrives in Step 2")


# ---------- Step 3: tables ----------
def parse_table(layout):
    """Read a table out of layout-mode text, where columns are separated by runs of 2+ spaces.
    The header is the first line with 3 or more cells; rows are the following lines with the same number of cells.
    Returns (rows as dicts keyed by header, every other non-empty line of the page, stripped)."""
    raise NotImplementedError("parse_table() arrives in Step 3")


def row_text(caption, row):
    """One self-describing line per row: "<caption>. Property type: Bedsit; Rent: £78.40; ..." in column order."""
    raise NotImplementedError("row_text() arrives in Step 3")


def layout_text(folder, doc, page):
    """The page's text with its layout kept (pypdf extraction_mode="layout")."""
    return pypdf.PdfReader(os.path.join(folder, doc)).pages[page - 1].extract_text(extraction_mode="layout")


# ---------- Step 4: scanned pages ----------
TRANSCRIBE_PROMPT = "Transcribe this scanned page exactly, line by line. Output only the text on the page."


def page_image(folder, doc, page):
    """The bytes of the first image on the page (a scan is one big JPEG)."""
    return pypdf.PdfReader(os.path.join(folder, doc)).pages[page - 1].images[0].data


def transcribe(image_bytes):
    """Read a scanned page with the vision model: one user message holding TRANSCRIBE_PROMPT and the image as a
    base64 data URL (data:image/jpeg;base64,...). Returns the transcription."""
    raise NotImplementedError("transcribe() arrives in Step 4")


# ---------- Step 5: chunks ----------
def title_of(doc):
    """"tenant-handbook.pdf" -> "Tenant Handbook"."""
    return doc[:-4].replace("-", " ").title()


def build_chunks(folder="pdfs"):
    """Every document as chunks {"doc", "kind", "text"}: a section of prose per heading, one chunk per table row plus
    one for the text around the table, one chunk for a scanned page. Each chunk's text starts with where it is from."""
    raise NotImplementedError("build_chunks() arrives in Step 5")