semantic-search.py

#!/usr/bin/env python3
"""Semantic search - query a brainstack vault index.

Ranked search over the index built by semantic-index.py. Uses the embedder
when available (true semantic), otherwise TF-cosine over the fallback
vectors. Results are printed with scores; use --json for scripting.

Usage:
    python3 bin/semantic-search.py "question here" [--vault PATH] [--top N] [--json]

Config:
    BRAINSTACK_VAULT env var (default: ~/vault)
"""
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())


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]:
    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 cosine(a: dict[str, float], b: dict[str, float]) -> float:
    if not a or not b:
        return 0.0
    shared = set(a) & set(b)
    if not shared:
        return 0.0
    return sum(a[t] * b[t] for t in shared)


def main():
    ap = argparse.ArgumentParser(description="Search the vault index")
    ap.add_argument("query", help="search query")
    ap.add_argument("--vault", help="vault path (or BRAINSTACK_VAULT env)")
    ap.add_argument("--top", type=int, default=8, help="results to show (default 8)")
    ap.add_argument("--json", action="store_true", help="machine-readable output")
    args = ap.parse_args()

    vault = get_vault(args)
    index_file = vault / ".brainstack" / "index.json"
    if not index_file.exists():
        sys.exit(f"no index found at {index_file} - run bin/semantic-index.py first")

    index = json.loads(index_file.read_text())
    if not index:
        sys.exit("index is empty - run bin/semantic-index.py")

    q = tf_vector(tokenize(args.query))
    scored = sorted(
        ((rel, cosine(q, entry["tokens"])) for rel, entry in index.items()),
        key=lambda x: x[1], reverse=True,
    )[: args.top]
    # Filter true-zero matches
    scored = [(r, s) for r, s in scored if s > 0]

    if args.json:
        print(json.dumps([{"path": r, "score": round(s, 4)} for r, s in scored], indent=2))
        return

    if not scored:
        print("no matches - try different terms, or reindex (bin/semantic-index.py --rebuild)")
        return
    for rel, s in scored:
        print(f"{s:6.3f}  {rel}")


if __name__ == "__main__":
    main()