S117 · RAG I — Chunking Strategies & Indexing

"""rag_ingest.py — chunk a Markdown corpus with recursive + structural splitters.
 
Requires: pip install langchain-text-splitters sentence-transformers chromadb
"""
from __future__ import annotations
import hashlib
import re
from dataclasses import dataclass, field
from pathlib import Path
from langchain_text_splitters import (
    RecursiveCharacterTextSplitter,
    MarkdownHeaderTextSplitter,
)
from sentence_transformers import SentenceTransformer
import chromadb
 
# ---------- Config ----------
CORPUS_DIR = Path("./corpus")           # any folder of .md files
CHUNK_SIZE = 512                        # target tokens
CHUNK_OVERLAP = 80                      # ~15% overlap
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"   # 22M, CPU-friendly
COLLECTION_NAME = "docs_v1"
 
 
@dataclass
class Chunk:
    id: str
    text: str
    metadata: dict = field(default_factory=dict)
 
 
def load_markdown(path: Path) -> list[dict]:
    """Split a Markdown file first by headings, then by size within each section."""
    md_text = path.read_text(encoding="utf-8")
    # Structural pass: split on H1/H2/H3 → sections with heading breadcrumb.
    md_splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=[("#", "h1"), ("##", "h2"), ("###", "h3")]
    )
    sections = md_splitter.split_text(md_text)
    # Size pass: recursive splitter inside each section.
    size_splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE * 4,        # rough chars ≈ 4× tokens for English
        chunk_overlap=CHUNK_OVERLAP * 4,
        separators=["\n\n", "\n", ". ", " ", ""],
    )
    docs = []
    for sec in sections:
        # Build a heading breadcrumb: 'H1 > H2 > H3'.
        breadcrumb = " > ".join(
            sec.metadata[k] for k in ("h1", "h2", "h3") if k in sec.metadata
        )
        for i, piece in enumerate(size_splitter.split_text(sec.page_content)):
            docs.append({
                "text": (f"{breadcrumb}\n\n{piece}" if breadcrumb else piece),
                "metadata": {
                    "source": str(path),
                    "breadcrumb": breadcrumb,
                    "chunk_index": i,
                },
            })
    return docs
 
 
def deterministic_id(source: str, text: str) -> str:
    """Stable id so re-ingesting the same content is a no-op (idempotent)."""
    h = hashlib.sha1(f"{source}\n{text}".encode()).hexdigest()
    return h[:16]
 
 
def build_index() -> None:
    files = sorted(CORPUS_DIR.rglob("*.md"))
    print(f"Loading {len(files)} markdown files from {CORPUS_DIR} ...")
 
    chunks: list[Chunk] = []
    for f in files:
        for d in load_markdown(f):
            chunks.append(Chunk(
                id=deterministic_id(d["metadata"]["source"], d["text"]),
                text=d["text"],
                metadata=d["metadata"],
            ))
    print(f"Produced {len(chunks)} chunks. Sample metadata:")
    for c in chunks[:3]:
        print(f"  [{c.id}] {c.metadata['breadcrumb'] or '(no heading)'}"
              f"  · {len(c.text)} chars")
 
    # Embed in batches. MiniLM outputs 384-dim vectors, ~100 chunks/sec on CPU.
    encoder = SentenceTransformer(EMBED_MODEL)
    texts = [c.text for c in chunks]
    print("\nEmbedding ...")
    vectors = encoder.encode(texts, show_progress_bar=True, batch_size=32)
 
    # Store in Chroma (persistent, on-disk).
    client = chromadb.PersistentClient(path="./chroma_store")
    coll = client.get_or_create_collection(name=COLLECTION_NAME)
    coll.upsert(
        ids=[c.id for c in chunks],
        documents=[c.text for c in chunks],
        metadatas=[c.metadata for c in chunks],
        embeddings=vectors.tolist(),
    )
    print(f"\nIndexed {len(chunks)} chunks into '{COLLECTION_NAME}'"
          f" ({coll.count()} total).")
 
    # Quick sanity query.
    demo_q = "How do I refresh an OAuth token?"
    q_vec = encoder.encode([demo_q])[0]
    hits = coll.query(query_embeddings=[q_vec.tolist()], n_results=3)
    print(f"\nSanity query: {demo_q!r}")
    for i, (doc, meta) in enumerate(zip(hits["documents"][0], hits["metadatas"][0])):
        print(f"  #{i+1}  {meta.get('breadcrumb','(no heading)')}")
        print(f"        {doc[:120].replace(chr(10),' ')}...")
 
 
if __name__ == "__main__":
    build_index()