ingest/parse_books.py

"""Turn the rendered markdown into retrieval chunks.

Reads `data/manifest.json` and the markdown under `data/sources/`, and writes one
JSON object per chunk to `data/chunks.jsonl`.

We slide a window down each page one line at a time. The window tracks what it
needs in order to describe itself — the headings above it, whether it is inside a
fence, which code it has taken in — so that we can read a chunk's metadata off it
at the moment it closes, rather than working it out again afterwards.

The window closes on a heading, or on a paragraph break once it is over budget,
and never inside a fence in order to keep code blocks whole.

    uv run python -m ingest.parse_books
    uv run python -m ingest.parse_books --book book --limit 5 --show
"""

import argparse
import hashlib
import json
import re
import sys
from collections import Counter
from dataclasses import dataclass, field

from transformers import AutoTokenizer

from rag.config import (
    BOOKS,
    CHUNK_HEADING_LEVEL,
    CHUNK_MAX_TOKENS,
    CHUNK_TARGET_TOKENS,
    CHUNKS_PATH,
    EMBED_MODEL,
    HEADING_SEPARATOR,
    MANIFEST_PATH,
    SOURCES_DIR,
)
from rag.types import Chunk, ManifestBook, Page

QUOTE = re.compile(r"^(\s*>\s?)+")
FENCE = re.compile(r"^\s*```(?P<info>[^`]*)$")
HEADING = re.compile(r"^(?P<hashes>#{1,6})\s+(?P<text>.+?)\s*$")

HIDDEN_LINE = re.compile(r"^#(\s|$)")
RUST_FENCE_TAGS = ("rust", "ignore", "should_panic", "no_run", "compile_fail", "edition")

RULE_MARKER = re.compile(r"^r\[[a-z0-9._-]+\]$")
LINK_DEFINITION = re.compile(r"^\[[^^\]]+\]:\s*\S+$")
LEGACY_ANCHOR = re.compile(r'^<a\s+id="[^"]*"\s*>\s*</a>$')
CALLOUT = re.compile(r"^\[!\w+\]\s*$")

LISTING_OPEN = re.compile(r"<Listing\b(?P<attrs>[^>]*)>")
LISTING_CLOSE = "</Listing>"
ATTRIBUTE = re.compile(r'(?P<key>[\w-]+)="(?P<value>[^"]*)"')

COMMENT_OPEN, COMMENT_CLOSE = "<!--", "-->"

INLINE_LINK = re.compile(r"!?\[(?P<label>[^\]]*)\]\((?:[^()]|\([^()]*\))*\)")
REFERENCE_LINK = re.compile(r"!?\[(?P<label>[^\]]*)\]\[[^\]]*\]")

EMPHASIS = re.compile(
    r"\*\*(.+?)\*\*|__(.+?)__|\*(.+?)\*|(?<![A-Za-z0-9])_(?=\S)(.+?)(?<=\S)_(?![A-Za-z0-9])"
)

FORBIDDEN_SUBSTRINGS = ("{{#", "ANCHOR:", "<Listing", COMMENT_OPEN)


@dataclass
class Window:
    """The chunk we are currently building, and what we can tell about it so far."""

    lines: list[str] = field(default_factory=list)
    tokens: int = 0
    code_tags: list[str] = field(default_factory=list)

    def add(self, line: str, cost: int) -> None:
        self.lines.append(line)
        self.tokens += cost

    def absorb(self, other: "Window") -> None:
        self.lines.extend(other.lines)
        self.code_tags.extend(other.code_tags)
        self.tokens += other.tokens
        other.reset()

    def reset(self) -> None:
        self.lines.clear()
        self.code_tags.clear()
        self.tokens = 0

    @property
    def body(self) -> str:
        return strip_links("\n".join(self.lines)).strip()


def strip_links(text: str) -> str:
    text = INLINE_LINK.sub(lambda m: m.group("label"), text)
    return REFERENCE_LINK.sub(lambda m: m.group("label"), text)


def strip_emphasis(text: str) -> str:
    return EMPHASIS.sub(lambda m: next(g for g in m.groups() if g is not None), text)


def slugify(text: str) -> str:
    """A heading's anchor id, matching mdBook's `normalize_id`."""
    text = strip_emphasis(strip_links(text))
    kept = "".join(c for c in text if c.isalnum() or c in "_- ")
    return kept.replace(" ", "-").lower()


def listing_caption(attrs: str) -> str:
    found = {m.group("key"): m.group("value") for m in ATTRIBUTE.finditer(attrs)}
    label = f"Listing {found['number']}" if "number" in found else "Listing"
    if found.get("caption"):
        label += f": {found['caption']}"
    if found.get("file-name"):
        label += f" ({found['file-name']})"
    return label


def is_markup(line: str) -> bool:
    return bool(
        RULE_MARKER.match(line)
        or LINK_DEFINITION.match(line)
        or LEGACY_ANCHOR.match(line)
        or CALLOUT.match(line)
    )


def chunk_page(page: Page, book_title: str, text: str, tokenizer) -> list[Chunk]:
    chunks: list[Chunk] = []
    window = Window()  # the chunk we are building
    unit = Window()  # the paragraph or fence we are currently reading
    headings: dict[int, tuple[str, str]] = {}
    anchors: Counter[str] = Counter()
    in_fence = False
    in_comment = False
    fence_info = ""

    def close() -> None:
        body, tags = window.body, sorted(set(window.code_tags))
        window.reset()
        if not body:
            return
        path = [headings[level][0] for level in sorted(headings)]
        anchor = next(
            (headings[level][1] for level in sorted(headings, reverse=True) if headings[level][1]),
            None,
        )
        breadcrumb = HEADING_SEPARATOR.join([book_title, *path])
        chunks.append(
            {
                "id": f"{page['book']}/{page['path']}#{anchor or 'section'}-{len(chunks)}",
                "text": f"{breadcrumb}\n\n{body}",
                "metadata": {
                    "book": page["book"],
                    "book_title": book_title,
                    "part": page["part"] or "",
                    "chapter": page["title"],
                    "heading_path": path,
                    "path": page["path"],
                    "url": page["url"] + (f"#{anchor}" if anchor else ""),
                    "has_code": bool(tags),
                    "code_tags": tags,
                },
            }
        )

    def commit() -> None:
        """Move the finished unit into the window, closing the window first if the
        unit would take it over budget."""
        if not unit.lines:
            return
        if window.lines and window.tokens + unit.tokens > CHUNK_TARGET_TOKENS:
            close()
        window.absorb(unit)

    def cost(line: str) -> int:
        return len(tokenizer.encode(line, add_special_tokens=False))

    for raw in text.splitlines():
        line = QUOTE.sub("", raw)

        fence = FENCE.match(line)
        if fence and not in_comment:
            unit.add(line.lstrip(), cost(line))
            if in_fence:
                in_fence = False
                commit()
            else:
                in_fence = True
                fence_info = fence.group("info").strip()
                unit.code_tags.append(fence_info or "(untagged)")
            continue

        if in_fence:
            # Only Rust fences, since `#` is an ordinary comment in most other
            # languages and stripping it would delete content.
            if any(tag in fence_info for tag in RUST_FENCE_TAGS):
                if HIDDEN_LINE.match(line):
                    continue
                line = re.sub(r"^##", "#", line)
            unit.add(line, cost(line))
            continue

        line, in_comment = strip_comments(line, in_comment)
        stripped = line.strip()

        heading = HEADING.match(stripped)
        if heading:
            commit()
            level = len(heading.group("hashes"))
            if level <= CHUNK_HEADING_LEVEL:
                close()
            base = slugify(heading.group("text"))
            anchors[base] += 1
            headings = {lvl: v for lvl, v in headings.items() if lvl < level}
            headings[level] = (
                strip_emphasis(strip_links(heading.group("text"))).strip(),
                base if anchors[base] == 1 else f"{base}-{anchors[base] - 1}",
            )
            continue

        listing = LISTING_OPEN.search(stripped)
        if listing:
            caption = listing_caption(listing.group("attrs"))
            unit.add(caption, cost(caption))
            continue
        if LISTING_CLOSE in stripped or is_markup(stripped):
            continue

        if not stripped:
            commit()
            continue

        unit.add(line, cost(line))

    commit()
    close()
    return chunks


def parse_book(
    book: str, manifest_book: ManifestBook, tokenizer, limit: int | None
) -> list[Chunk]:
    chunks: list[Chunk] = []
    pages = manifest_book["pages"]
    if limit:
        pages = pages[:limit]

    for page in pages:
        source = SOURCES_DIR / book / page["path"]
        if not source.exists():
            print(f"  missing: {page['path']}", file=sys.stderr)
            continue
        chunks.extend(
            chunk_page(page, manifest_book["title"], source.read_text(encoding="utf-8"), tokenizer)
        )

    return chunks


def strip_comments(line: str, in_comment: bool) -> tuple[str, bool]:
    """Remove HTML comment spans, carrying the open state to the next line."""
    kept: list[str] = []
    while line:
        if in_comment:
            _, found, line = line.partition(COMMENT_CLOSE)
            if not found:
                return "".join(kept), True
            in_comment = False
        else:
            before, found, line = line.partition(COMMENT_OPEN)
            kept.append(before)
            if not found:
                break
            in_comment = True
    return "".join(kept), in_comment


def hidden_lines(text: str) -> list[str]:
    found, info = [], None
    for line in text.split("\n"):
        fence = FENCE.match(line)
        if fence:
            info = None if info is not None else fence.group("info")
            continue
        if info is not None and any(t in info for t in RUST_FENCE_TAGS) and HIDDEN_LINE.match(line):
            found.append(line)
    return found


def check(chunks: list[Chunk]) -> None:
    unbalanced = [c["id"] for c in chunks if c["text"].count("```") % 2]
    if unbalanced:
        raise ValueError(f"{len(unbalanced)} chunks split a code block, e.g. {unbalanced[0]}")

    for needle in FORBIDDEN_SUBSTRINGS:
        hits = [c["id"] for c in chunks if needle in c["text"]]
        if hits:
            raise ValueError(f"residue {needle!r} left in {len(hits)} chunks, e.g. {hits[0]}")

    empty = [c["id"] for c in chunks if not c["text"].split("\n\n", 1)[-1].strip()]
    if empty:
        raise ValueError(f"{len(empty)} chunks have a breadcrumb but no body, e.g. {empty[0]}")

    # the books render hidden lines in the code blocks which we need to ensure have been stripped out of the corpus otherwise they may surface in the app as noise in the answer
    for chunk in chunks:
        left = hidden_lines(chunk["text"])
        if left:
            raise ValueError(f"rustdoc-hidden line survived in {chunk['id']}: {left[0]!r}")


def report(chunks: list[Chunk], tokenizer) -> None:
    lengths = [len(tokenizer.encode(c["text"], add_special_tokens=False)) for c in chunks]
    counts: Counter[str] = Counter()
    tokens: Counter[str] = Counter()
    for chunk, length in zip(chunks, lengths):
        counts[chunk["metadata"]["book"]] += 1
        tokens[chunk["metadata"]["book"]] += length

    print("\nchunks per book:")
    for book in sorted(counts):
        print(f"   {book:18} {counts[book]:>5} chunks  {tokens[book]:>8,} tokens")

    ordered = sorted(lengths)
    over = sum(1 for n in ordered if n > CHUNK_MAX_TOKENS)
    print(f"\ntotal: {len(chunks):,} chunks, {sum(ordered):,} tokens")
    print(f"   median {ordered[len(ordered) // 2]}, max {ordered[-1]}, {over} over {CHUNK_MAX_TOKENS}")


def deduplicate(chunks: list[Chunk]) -> tuple[list[Chunk], int]:
    seen: set[str] = set()
    kept: list[Chunk] = []
    for chunk in chunks:
        digest = hashlib.sha256(chunk["text"].encode()).hexdigest()
        if digest in seen:
            continue
        seen.add(digest)
        kept.append(chunk)
    return kept, len(chunks) - len(kept)


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--book", action="append", choices=sorted(BOOKS))
    parser.add_argument("--limit", type=int, help="only the first N pages per book")
    parser.add_argument("--show", action="store_true", help="print the first chunk of each book")
    args = parser.parse_args(argv[1:])

    if not MANIFEST_PATH.exists():
        print("No manifest. Run: uv run python -m ingest.fetch_sources", file=sys.stderr)
        return 2

    manifest = json.loads(MANIFEST_PATH.read_text())
    tokenizer = AutoTokenizer.from_pretrained(EMBED_MODEL)

    chunks: list[Chunk] = []
    for book in args.book or sorted(manifest["books"]):
        print(f"parsing {book}...", flush=True)
        book_chunks = parse_book(book, manifest["books"][book], tokenizer, args.limit)
        if args.show and book_chunks:
            print("-" * 70)
            print(book_chunks[0]["text"][:700])
            print("-" * 70)
        chunks.extend(book_chunks)

    partial = bool(args.book or args.limit)
    chunks, duplicates = deduplicate(chunks)
    if duplicates:
        print(f"\ncollapsed {duplicates} exact-duplicate chunks")

    check(chunks)
    report(chunks, tokenizer)

    if partial:
        # A --book or --limit run is for inspection only. Writing it would replace
        # the whole corpus with a fragment, which the retriever and BM25 would then
        # load as if it were complete.
        print(f"\npartial run ({len(chunks)} chunks): {CHUNKS_PATH.name} left untouched")
        return 0

    CHUNKS_PATH.parent.mkdir(parents=True, exist_ok=True)
    with CHUNKS_PATH.open("w", encoding="utf-8") as handle:
        for chunk in chunks:
            handle.write(json.dumps(chunk, ensure_ascii=False) + "\n")
    print(f"\nwrote {CHUNKS_PATH.name} ({CHUNKS_PATH.stat().st_size / 1e6:.1f} MB)")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))