scripts/build_index.py

"""Chunk the Zephyr docs and build a FAISS index over them.

Chunking is the part that decides whether retrieval works, so it is done on
document structure rather than on a fixed character count. RST carries its own
section headings; splitting on those keeps a chunk to one topic and lets every
retrieved passage cite the section it came from. A blind 1000-character split
would cut mid-sentence and mix two subsystems into one vector.

    python scripts/build_index.py
    python scripts/build_index.py --docs data/raw_docs --out data/index
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_DOCS = ROOT / "data" / "raw_docs"
DEFAULT_OUT = ROOT / "data" / "index"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

# An RST section underline: a run of one punctuation character on its own line,
# directly under the title it underlines.
RST_UNDERLINE = re.compile(r"^([=\-`:'\"~^_*+#<>])\1{2,}\s*$")

MIN_CHARS = 120
MAX_CHARS = 1600

# Files that live under doc/ but are not documentation. Indexing them puts a
# pip requirements file and a JavaScript 404 page into the retrieval pool.
SKIP_FILES = {
    "404.rst",
    "CMakeLists.txt",
    "requirements.txt",
    "substitutions.txt",
    "index.rst",
}

# Release notes end with thousands of bare GitHub issue IDs. They are the
# largest sections in the corpus and answer nothing: one release note produced a
# single 104 KB chunk of numbers.
SKIP_SECTIONS = re.compile(r"issue related items|^bugs?$|security vulnerability", re.IGNORECASE)

# A chunk that is mostly punctuation, IDs or markup is not prose. Retrieval will
# happily return it and the model will have nothing to say about it.
def looks_like_prose(text: str) -> bool:
    letters = sum(character.isalpha() for character in text)
    if letters < len(text) * 0.45:
        return False
    words = text.split()
    if not words:
        return False
    # Long runs of short tokens are ID lists and option tables, not sentences.
    return sum(len(word) for word in words) / len(words) >= 3.0


def clean(text: str) -> str:
    """Strip the RST directives that carry no meaning for a reader."""
    lines: list[str] = []
    for line in text.splitlines():
        stripped = line.strip()
        # Comments, toctrees and figure/image directives are navigation and
        # layout. They retrieve badly and answer nothing.
        if stripped.startswith(".. toctree::") or stripped.startswith(".. figure::"):
            continue
        if stripped.startswith(".. image::") or stripped.startswith(".."):
            if re.match(r"^\.\.\s+_[\w.-]+:", stripped):
                continue  # anchor target
            if stripped.startswith(".. code-block::") or stripped.startswith(".. note::"):
                lines.append(line)  # keep — the body that follows is content
                continue
            continue
        if stripped.startswith(":") and stripped.count(":") >= 2 and len(stripped) < 80:
            continue  # field list / directive option
        lines.append(line)
    return "\n".join(lines)


def sections(text: str) -> list[tuple[str, str]]:
    """Split RST into (heading, body). The first block inherits the document title."""
    lines = text.splitlines()
    blocks: list[tuple[str, list[str]]] = [("", [])]

    index = 0
    while index < len(lines):
        line = lines[index]
        following = lines[index + 1] if index + 1 < len(lines) else ""
        is_heading = (
            line.strip()
            and RST_UNDERLINE.match(following)
            and len(following.strip()) >= len(line.strip()) - 2
        )
        if is_heading:
            blocks.append((line.strip(), []))
            index += 2
            continue
        blocks[-1][1].append(line)
        index += 1

    return [(heading, "\n".join(body).strip()) for heading, body in blocks]


def split_long(body: str, limit: int = MAX_CHARS) -> list[str]:
    """Break an over-long section on blank lines, never mid-paragraph.

    Paragraph splitting alone is not enough. A section with no blank lines — a
    long table, a generated list — comes back as one piece however large it is.
    The first version of this used only blank lines and emitted a single
    126 KB chunk, so anything still over the limit is hard-split on line
    boundaries as a backstop.
    """
    if len(body) <= limit:
        return [body]

    parts: list[str] = []
    current = ""
    for paragraph in body.split("\n\n"):
        if current and len(current) + len(paragraph) + 2 > limit:
            parts.append(current.strip())
            current = paragraph
        else:
            current = f"{current}\n\n{paragraph}" if current else paragraph
    if current.strip():
        parts.append(current.strip())

    bounded: list[str] = []
    for part in parts:
        if len(part) <= limit:
            bounded.append(part)
            continue
        buffer = ""
        for line in part.splitlines():
            if buffer and len(buffer) + len(line) + 1 > limit:
                bounded.append(buffer.strip())
                buffer = line
            else:
                buffer = f"{buffer}\n{line}" if buffer else line
        if buffer.strip():
            bounded.append(buffer.strip())
    return bounded


def chunk_document(path: Path) -> list[dict]:
    raw = path.read_text(encoding="utf-8", errors="replace")
    # The fetcher flattens "kernel/services/threads.rst" to
    # "kernel__services__threads.rst" so the source path survives as a citation.
    source = path.stem.replace("__", "/")
    body = clean(raw)

    title = ""
    for heading, _ in sections(body):
        if heading:
            title = heading
            break

    chunks: list[dict] = []
    for heading, section_body in sections(body):
        if len(section_body) < MIN_CHARS:
            continue
        if heading and SKIP_SECTIONS.search(heading):
            continue
        for part in split_long(section_body):
            if len(part) < MIN_CHARS or not looks_like_prose(part):
                continue
            chunks.append(
                {
                    "text": part,
                    "source": source,
                    "title": title or source,
                    "section": heading or title or source,
                    "url": f"https://docs.zephyrproject.org/latest/{source}.html",
                }
            )
    return chunks


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--docs", type=Path, default=DEFAULT_DOCS)
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--model", default=EMBEDDING_MODEL)
    parser.add_argument("--batch-size", type=int, default=64)
    args = parser.parse_args()

    if not args.docs.is_dir():
        raise SystemExit(f"no docs at {args.docs} — run scripts/fetch_docs.py first")

    files = sorted(p for p in args.docs.glob("*") if p.suffix.lower() in {".rst", ".md", ".txt"})
    files = [p for p in files if p.name != "_SOURCE.txt" and p.name not in SKIP_FILES]
    print(f"Chunking {len(files)} documents...")

    chunks: list[dict] = []
    for path in files:
        chunks.extend(chunk_document(path))
    if not chunks:
        raise SystemExit("no chunks produced — check the docs directory")

    lengths = sorted(len(c["text"]) for c in chunks)
    print(
        f"{len(chunks)} chunks | median {lengths[len(lengths) // 2]} chars, "
        f"max {lengths[-1]}"
    )

    # Imported here so --help and the chunking stats stay fast on a machine
    # without torch installed.
    import faiss
    import numpy as np
    from sentence_transformers import SentenceTransformer

    print(f"Embedding with {args.model}...")
    model = SentenceTransformer(args.model)
    vectors = model.encode(
        [c["text"] for c in chunks],
        batch_size=args.batch_size,
        show_progress_bar=True,
        convert_to_numpy=True,
        normalize_embeddings=True,
    ).astype("float32")

    # Inner product over L2-normalised vectors is cosine similarity, which is
    # what the retrieval scores in ask.py are reported as.
    index = faiss.IndexFlatIP(vectors.shape[1])
    index.add(vectors)

    args.out.mkdir(parents=True, exist_ok=True)
    faiss.write_index(index, str(args.out / "docs.faiss"))
    with (args.out / "chunks.jsonl").open("w", encoding="utf-8") as handle:
        for chunk in chunks:
            handle.write(json.dumps(chunk, ensure_ascii=False) + "\n")

    source_note = (args.docs / "_SOURCE.txt")
    meta = {
        "embedding_model": args.model,
        "dimensions": int(vectors.shape[1]),
        "chunks": len(chunks),
        "documents": len(files),
        "index": "IndexFlatIP (cosine over normalised vectors)",
        "source": source_note.read_text(encoding="utf-8") if source_note.exists() else "unknown",
    }
    (args.out / "meta.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")

    print(f"\nIndex written to {args.out}")
    print(f"{len(chunks)} chunks x {vectors.shape[1]} dims")
    return 0


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