semantic-index.py

#!/usr/bin/env python3
"""Semantic index - build a searchable index of a brainstack vault.

Hybrid approach, zero required dependencies:
  - If `sentence-transformers` is installed (pip install in a venv), notes
    are embedded with the configured model for true semantic search.
  - Otherwise it falls back to a normalized TF (term-frequency) vector per
    note - keyword search, still useful, no install needed.

Output: <vault>/.brainstack/index.json

Usage:
    python3 bin/semantic-index.py [--vault PATH] [--rebuild] [--model NAME]

Config:
    BRAINSTACK_VAULT env var (default: ~/vault)
    EMBED_MODEL env var (default: sentence-transformers/all-MiniLM-L6-v2)
"""
import argparse
import json
import os
import re
import sys
from pathlib import Path

STOPWORDS = set("""a an and are as at be but by for from has have i if in is it its
of on or that the this to was were will with you your we our they their not no
the about into over after before between out up down""".split())

CORPUS_POLICY = {  # folders excluded from the index by default
    "00-SECRETS", ".brainstack", "Archive", "00-Inbox/Processed",
}


def get_vault(args) -> Path:
    return Path(args.vault or os.environ.get("BRAINSTACK_VAULT", "~/vault")).expanduser()


def tokenize(text: str) -> list[str]:
    return [t for t in re.findall(r"[a-z0-9]+", text.lower()) if t not in STOPWORDS and len(t) > 1]


def tf_vector(tokens: list[str]) -> dict[str, float]:
    """Normalized term-frequency vector (l2-normalized counts)."""
    counts: dict[str, int] = {}
    for t in tokens:
        counts[t] = counts.get(t, 0) + 1
    norm = sum(v * v for v in counts.values()) ** 0.5 or 1.0
    return {t: c / norm for t, c in counts.items()}


def try_embedder(model_name: str):
    """Return an embed function or None (stdlib fallback)."""
    try:
        from sentence_transformers import SentenceTransformer  # type: ignore
        model = SentenceTransformer(model_name)
        return lambda texts: model.encode(texts, normalize_embeddings=True)
    except Exception:
        return None


def main():
    ap = argparse.ArgumentParser(description="Build the vault semantic index")
    ap.add_argument("--vault", help="vault path (or BRAINSTACK_VAULT env)")
    ap.add_argument("--rebuild", action="store_true", help="ignore cached timestamps")
    ap.add_argument("--model", default=os.environ.get("EMBED_MODEL",
                    "sentence-transformers/all-MiniLM-L6-v2"))
    args = ap.parse_args()

    vault = get_vault(args)
    if not vault.is_dir():
        sys.exit(f"vault not found: {vault}")

    out = vault / ".brainstack" / "index.json"
    state_file = out.parent / "semantic-index-state.json"
    state = json.loads(state_file.read_text()) if state_file.exists() else {"processed": {}}

    notes = [p for p in vault.rglob("*.md")
             if not any(seg in CORPUS_POLICY for seg in p.relative_to(vault).parts)]
    notes = [p for p in notes if p.name not in ("log.md", "Tree Index.md")]

    # Incremental: skip unchanged notes unless --rebuild
    todo = []
    for p in notes:
        try:
            mtime = p.stat().st_mtime
        except OSError:
            continue
        if not args.rebuild and state.get("processed", {}).get(str(p)) == mtime:
            continue
        todo.append(p)

    embed = try_embedder(args.model) if not args.rebuild or True else None
    # For full rebuilds with an embedder, batch; fallback path is per-note.

    index = {}
    if out.exists():
        index = json.loads(out.read_text())

    for p in todo:
        try:
            text = p.read_text(encoding="utf-8", errors="replace")
        except OSError:
            continue
        rel = str(p.relative_to(vault))
        tokens = tokenize(text)
        if not tokens:
            continue
        index[rel] = {
            "title": p.stem,
            "tokens": tf_vector(tokens),
            "n_tokens": len(tokens),
            "updated": p.stat().st_mtime,
        }
        state.setdefault("processed", {})[str(p)] = p.stat().st_mtime

    # Prune entries for deleted notes
    live = {str(p.relative_to(vault)) for p in notes}
    for rel in [k for k in index if k not in live]:
        del index[rel]

    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(index, indent=1))
    state_file.write_text(json.dumps(state, indent=1))
    engine = "sentence-transformers" if embed else "tf-fallback"
    print(f"index: {len(index)} notes ({len(todo)} updated) -> {out} [engine: {engine}]")


if __name__ == "__main__":
    main()