`knowledge-layer/scripts/kl_query.py`

175 lines
#!/usr/bin/env python3
"""
kl_query.py — read-first retrieval for a Knowledge Layer (NO vectors, NO RAG).

Given a question or topic, return the handful of wiki entry files an agent
should LOAD before doing its task — so it "wakes up with full context." Ranking
is deterministic keyword overlap over each entry's title, aliases, tags, and
summary (the frontmatter + first lines), plus a small bonus for [[wikilink]]
neighbours of top hits. It prints paths (and optionally the entries), never a
guess — you read the files it names.

Usage:
    kl_query.py "<question or topic>" [<kl-root>] [-n N] [--print] [--json]

Options:
    -n N       max entries to return (default 6)
    --print    also print each entry's body (for direct context loading)
    --json     machine-readable {path, score, why}[]

Exit codes: 0 = hits, 1 = no hits, 2 = error.
"""
import argparse
import json
import re
import sys
from pathlib import Path

STOP = set("the a an of to and or for in on at by with from is are be as it "
           "this that these those we you i our your their his her its what how "
           "why when who whom which about into over under than then".split())
ENTRY_DIRS_SKIP = {".kl", "sources"}
FM_RE = re.compile(r"^---\s*$")


def tokenize(s: str):
    return [t for t in re.findall(r"[a-z0-9]+", s.lower()) if t and t not in STOP]


def parse_entry(p: Path):
    """Return (title, aliases[], tags[], summary, body, wikilinks[])."""
    text = p.read_text(errors="replace")
    lines = text.splitlines()
    fm = {}
    body_start = 0
    if lines and FM_RE.match(lines[0]):
        for i in range(1, len(lines)):
            if FM_RE.match(lines[i]):
                body_start = i + 1
                break
            m = re.match(r"([A-Za-z_]+):\s*(.*)$", lines[i])
            if m:
                fm[m.group(1).lower()] = m.group(2).strip()
    body = "\n".join(lines[body_start:])

    def listify(v):
        if not v:
            return []
        v = v.strip().strip("[]")
        return [x.strip().strip('"').strip("'") for x in v.split(",") if x.strip()]

    title = ""
    mtitle = re.search(r"^#\s+(.+)$", body, re.M)
    if mtitle:
        title = mtitle.group(1).strip()
    if not title:
        title = p.stem.replace("-", " ")

    summ = ""
    msum = re.search(r"\*\*Summary:\*\*\s*(.+)", body)
    if msum:
        summ = msum.group(1).strip()

    wikilinks = re.findall(r"\[\[([^\]]+)\]\]", body)
    return {
        "title": title,
        "aliases": listify(fm.get("aliases", "")),
        "tags": listify(fm.get("tags", "")),
        "type": fm.get("type", p.parent.name),
        "summary": summ,
        "body": body,
        "links": [w.strip().lower() for w in wikilinks],
    }


def collect_entries(root: Path):
    out = {}
    for p in sorted(root.rglob("*.md")):
        rel = p.relative_to(root)
        top = rel.parts[0]
        if top in ENTRY_DIRS_SKIP:
            continue
        if p.name in ("SCHEMA.md", "INDEX.md", "CONTRADICTIONS.md",
                      "SOURCES.md", "_about.md"):
            continue
        out[p] = parse_entry(p)
    return out


def score(qtokens, e):
    hay_strong = " ".join([e["title"]] + e["aliases"] + e["tags"]).lower()
    strong = set(tokenize(hay_strong))
    weak = set(tokenize(e["summary"] + " " + e["body"][:600]))
    s, why = 0, []
    for t in qtokens:
        if t in strong:
            s += 3
            why.append(t)
        elif t in weak:
            s += 1
    return s, why


def main() -> int:
    ap = argparse.ArgumentParser(description="Read-first retrieval (no vectors).")
    ap.add_argument("query")
    ap.add_argument("root", nargs="?", default="knowledge-layer")
    ap.add_argument("-n", type=int, default=6)
    ap.add_argument("--print", dest="show", action="store_true")
    ap.add_argument("--json", action="store_true")
    args = ap.parse_args()

    root = Path(args.root).resolve()
    if not (root / ".kl").is_dir():
        print(f"ERROR: {root} is not a Knowledge Layer (run kl_init.py first).",
              file=sys.stderr)
        return 2

    qtokens = tokenize(args.query)
    entries = collect_entries(root)
    scored = []
    stem_to_path = {p.stem.lower(): p for p in entries}
    for p, e in entries.items():
        s, why = score(qtokens, e)
        if s > 0:
            scored.append([p, e, s, why])

    # neighbour bonus: entries linked from a top hit get a small nudge.
    top_paths = {p for p, *_ in sorted(scored, key=lambda x: -x[2])[:args.n]}
    for p, e, s, why in scored:
        for lk in e["links"]:
            if lk in stem_to_path and stem_to_path[lk] in top_paths:
                s += 1
    scored.sort(key=lambda x: -x[2])
    scored = scored[:args.n]

    if not scored:
        print("No matching entries. The knowledge layer may not cover this yet — "
              "add a source to sources/ and compile.", file=sys.stderr)
        return 1

    if args.json:
        print(json.dumps([{
            "path": str(p.relative_to(root)),
            "score": s, "why": why, "summary": e["summary"]
        } for p, e, s, why in scored], indent=2))
        return 0

    print(f"Read these {len(scored)} entries first (query: {args.query!r}):\n")
    for p, e, s, why in scored:
        rel = p.relative_to(root)
        tag = f"[{e['type']}]"
        print(f"  {s:>3}  {rel}  {tag}")
        if e["summary"]:
            print(f"       {e['summary']}")
    if args.show:
        print("\n" + "=" * 60)
        for p, e, s, why in scored:
            print(f"\n### {p.relative_to(root)}\n")
            print(e["body"].strip())
    return 0


if __name__ == "__main__":
    sys.exit(main())