From 691d2d94f568069bd6c798766f27696b93fe0ced Mon Sep 17 00:00:00 2001
From: marcuspaico <marcus@thepai.co>
Date: Fri, 14 Aug 2026 17:18:42 -0700
Subject: [PATCH] feat: chunked bge embeddings + lancedb sutta index

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 scripts/build_index.py | 10 ++++++++++
 src/buddhagpt/index.py | 40 ++++++++++++++++++++++++++++++++++++++++
 tests/test_index.py    | 14 ++++++++++++++
 3 files changed, 64 insertions(+)
 create mode 100644 scripts/build_index.py
 create mode 100644 src/buddhagpt/index.py
 create mode 100644 tests/test_index.py

diff --git a/scripts/build_index.py b/scripts/build_index.py
new file mode 100644
index 0000000..ea1ebde
--- /dev/null
+++ b/scripts/build_index.py
@@ -0,0 +1,10 @@
+#!/usr/bin/env python
+"""Build the sutta embeddings index."""
+from pathlib import Path
+from buddhagpt.index import build_index
+
+if __name__ == "__main__":
+    jsonl = Path("corpus/suttas.jsonl")
+    db_path = Path("data/lancedb")
+    n = build_index(jsonl, db_path)
+    print(f"Built index with {n} chunks at {db_path}")
diff --git a/src/buddhagpt/index.py b/src/buddhagpt/index.py
new file mode 100644
index 0000000..19a5308
--- /dev/null
+++ b/src/buddhagpt/index.py
@@ -0,0 +1,40 @@
+import json
+from pathlib import Path
+import lancedb
+from sentence_transformers import SentenceTransformer
+
+EMBED_MODEL = "BAAI/bge-small-en-v1.5"
+
+
+def chunk_text(text: str, max_words: int = 300, overlap: int = 50) -> list[str]:
+    words = text.split()
+    chunks, start = [], 0
+    while start < len(words):
+        chunks.append(" ".join(words[start:start + max_words]))
+        if start + max_words >= len(words):
+            break
+        start += max_words - overlap
+    return chunks
+
+
+def build_index(jsonl: Path, db_path: Path) -> int:
+    model = SentenceTransformer(EMBED_MODEL, device="mps")
+    rows = []
+    for line in jsonl.read_text().splitlines():
+        s = json.loads(line)
+        for i, chunk in enumerate(chunk_text(s["text"])):
+            rows.append({"uid": s["uid"], "title": s["title"], "chunk": chunk, "chunk_i": i})
+    vecs = model.encode([r["chunk"] for r in rows], batch_size=64, show_progress_bar=True)
+    for r, v in zip(rows, vecs):
+        r["vector"] = v.tolist()
+    db = lancedb.connect(db_path)
+    db.create_table("suttas", rows, mode="overwrite")
+    return len(rows)
+
+
+def search(db_path: Path, query: str, k: int = 4) -> list[dict]:
+    model = SentenceTransformer(EMBED_MODEL, device="mps")
+    tbl = lancedb.connect(db_path).open_table("suttas")
+    q = model.encode([query])[0].tolist()
+    hits = tbl.search(q).limit(k).to_list()
+    return [{"uid": h["uid"], "title": h["title"], "chunk": h["chunk"], "score": h["_distance"]} for h in hits]
diff --git a/tests/test_index.py b/tests/test_index.py
new file mode 100644
index 0000000..38dc6df
--- /dev/null
+++ b/tests/test_index.py
@@ -0,0 +1,14 @@
+from buddhagpt.index import chunk_text
+
+
+def test_chunk_text_respects_max_words():
+    text = " ".join(f"w{i}" for i in range(700))
+    chunks = chunk_text(text, max_words=300, overlap=50)
+    assert all(len(c.split()) <= 300 for c in chunks)
+    assert len(chunks) == 3
+
+
+def test_chunks_overlap():
+    text = " ".join(f"w{i}" for i in range(400))
+    a, b = chunk_text(text, max_words=300, overlap=50)
+    assert a.split()[-50:] == b.split()[:50]