scripts/build_index.py

"""
SUPERSEDED by scripts/build_index_from_crawler.py — kept for reference only.
See README.md for the current build path.

Build the RAG vector index from local books (Appendix A) + crawled articles
(scripts/crawl_sources.py output). Runs entirely on the user's machine.

The output index (rag_index/) intentionally stores ONLY opaque chunk IDs and
raw text — no title, author, URL, or domain. The Space can therefore only
ever cite retrieved context as "Source 1", "Source 2", etc.

Usage:
    python build_index.py --books-dir ../books --crawled-dir ../crawled --out ../rag_index
    python build_index.py --push-dataset --dataset-repo-id yourname/dm-rag-index --hf-token hf_...
"""

from __future__ import annotations

import argparse
import hashlib
import json
from datetime import datetime, timezone
from pathlib import Path

import tiktoken
from bs4 import BeautifulSoup
from pypdf import PdfReader
from tqdm import tqdm

ENCODING = tiktoken.get_encoding("cl100k_base")
CHUNK_TOKENS = 800
OVERLAP_TOKENS = 100
EMBED_MODEL = "BAAI/bge-small-en-v1.5"
EMBED_BATCH = 64


def read_txt(path: Path) -> str:
    return path.read_text(encoding="utf-8", errors="ignore")


def read_pdf(path: Path) -> str:
    reader = PdfReader(str(path))
    return "\n".join(page.extract_text() or "" for page in reader.pages)


def read_epub(path: Path) -> str:
    import ebooklib
    from ebooklib import epub

    book = epub.read_epub(str(path))
    parts = []
    for item in book.get_items():
        if item.get_type() == ebooklib.ITEM_DOCUMENT:
            soup = BeautifulSoup(item.get_content(), "html.parser")
            parts.append(soup.get_text(separator="\n"))
    return "\n".join(parts)


READERS = {".txt": read_txt, ".pdf": read_pdf, ".epub": read_epub}


def chunk_text(text: str) -> list[str]:
    tokens = ENCODING.encode(text)
    if not tokens:
        return []
    chunks = []
    step = CHUNK_TOKENS - OVERLAP_TOKENS
    for start in range(0, len(tokens), step):
        window = tokens[start : start + CHUNK_TOKENS]
        if len(window) < 50:  # drop tiny tail fragments
            break
        chunks.append(ENCODING.decode(window))
        if start + CHUNK_TOKENS >= len(tokens):
            break
    return chunks


def opaque_id(text: str) -> str:
    return hashlib.sha1(text.encode("utf-8")).hexdigest()[:16]


def iter_book_documents(books_dir: Path):
    if not books_dir.exists():
        return
    for path in sorted(books_dir.iterdir()):
        reader = READERS.get(path.suffix.lower())
        if reader is None:
            continue
        try:
            text = reader(path)
        except Exception as exc:
            print(f"[book] failed to read {path.name}: {exc}")
            continue
        if text.strip():
            yield text


def iter_crawled_documents(crawled_dir: Path):
    if not crawled_dir.exists():
        return
    for txt_path in sorted(crawled_dir.glob("*/pages/*.txt")):
        text = read_txt(txt_path)
        if text.strip():
            yield text


def build_chunks(books_dir: Path, crawled_dir: Path) -> list[str]:
    all_chunks: list[str] = []
    seen_ids: set[str] = set()

    sources = [
        ("book", iter_book_documents(books_dir)),
        ("article", iter_crawled_documents(crawled_dir)),
    ]
    for kind, docs in sources:
        for doc_text in tqdm(list(docs), desc=f"chunking {kind}s"):
            for chunk in chunk_text(doc_text):
                cid = opaque_id(chunk)
                if cid in seen_ids:
                    continue  # de-dup identical chunks (e.g. re-run)
                seen_ids.add(cid)
                all_chunks.append(chunk)
    return all_chunks


def embed_chunks(chunks: list[str]):
    from sentence_transformers import SentenceTransformer

    model = SentenceTransformer(EMBED_MODEL)
    embeddings = model.encode(
        chunks, batch_size=EMBED_BATCH, show_progress_bar=True, normalize_embeddings=True
    )
    return embeddings


def write_index(chunks: list[str], embeddings, out_dir: Path):
    import chromadb

    out_dir.mkdir(parents=True, exist_ok=True)
    client = chromadb.PersistentClient(path=str(out_dir))
    collection = client.get_or_create_collection(
        name="dm_rag", metadata={"hnsw:space": "cosine"}
    )

    ids = [opaque_id(c) for c in chunks]
    for i in tqdm(range(0, len(chunks), EMBED_BATCH), desc="writing index"):
        batch_ids = ids[i : i + EMBED_BATCH]
        batch_docs = chunks[i : i + EMBED_BATCH]
        batch_emb = embeddings[i : i + EMBED_BATCH].tolist()
        # No metadatas passed: the index must never carry title/author/URL/domain.
        collection.upsert(ids=batch_ids, documents=batch_docs, embeddings=batch_emb)

    meta = {
        "model": EMBED_MODEL,
        "chunk_tokens": CHUNK_TOKENS,
        "overlap_tokens": OVERLAP_TOKENS,
        "count": len(chunks),
        "built_at": datetime.now(timezone.utc).isoformat(),
    }
    (out_dir / "index_meta.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
    print(f"Wrote {len(chunks)} chunks to {out_dir}")


def push_to_dataset(out_dir: Path, repo_id: str, hf_token: str):
    from huggingface_hub import HfApi

    api = HfApi(token=hf_token)
    api.create_repo(repo_id=repo_id, repo_type="dataset", private=True, exist_ok=True)
    api.upload_folder(folder_path=str(out_dir), repo_id=repo_id, repo_type="dataset")
    print(f"Pushed index to private dataset: {repo_id}")


def main():
    parser = argparse.ArgumentParser(description="Build the local RAG index")
    parser.add_argument("--books-dir", default="../books")
    parser.add_argument("--crawled-dir", default="../crawled")
    parser.add_argument("--out", default="../rag_index")
    parser.add_argument("--push-dataset", action="store_true")
    parser.add_argument("--dataset-repo-id", default=None)
    parser.add_argument("--hf-token", default=None)
    args = parser.parse_args()

    books_dir = Path(args.books_dir)
    crawled_dir = Path(args.crawled_dir)
    out_dir = Path(args.out)

    chunks = build_chunks(books_dir, crawled_dir)
    if not chunks:
        print("No chunks found — check --books-dir and --crawled-dir paths.")
        return

    embeddings = embed_chunks(chunks)
    write_index(chunks, embeddings, out_dir)

    if args.push_dataset:
        if not args.dataset_repo_id or not args.hf_token:
            raise SystemExit("--push-dataset requires --dataset-repo-id and --hf-token")
        push_to_dataset(out_dir, args.dataset_repo_id, args.hf_token)


if __name__ == "__main__":
    main()