eSlider
4f17b8e5ab
--with-facts / --facts-json write root=facts leafs; --with-chats picks up var/chats/md as info. WhatsApp sync stays out of v1 (Gitea #18).
243 lines
8.7 KiB
Python
Executable File
243 lines
8.7 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""kb/index - build the 2dph brain from markdown + factual leafs.
|
||
|
||
bin/kb/index [--corpus DIR] [--rebuild] [--limit N]
|
||
bin/kb/index --rebuild --with-facts --with-chats
|
||
bin/kb/index --json # emit stats as JSON
|
||
|
||
Reads every .md under the corpus (default: repo root docs, skills, READMEs)
|
||
as `info` leafs, embeds them with model2vec (potion-multilingual-128M), and
|
||
writes them into var/kb.lbug with FTS + HNSW indexes. `facts` leafs come
|
||
from bin/facts/extract (docker × compose × ssh-config pairing) when
|
||
`--with-facts` is set. `--with-chats` indexes markdown under var/chats/md
|
||
(or a given dir) as info. WhatsApp sync stays out of v1.
|
||
|
||
--rebuild drops the database file and indexes from scratch. Without it a run
|
||
is idempotent (MERGE by (source,text) id).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
ROOT = Path(__file__).resolve().parents[2]
|
||
sys.path.insert(0, str(ROOT / "bin" / "tools"))
|
||
|
||
from kblib import ( # noqa: E402
|
||
add_leafs, connect, ensure_indexes, init_schema, upsert_leaf, link_from_file,
|
||
open_readonly, stats,
|
||
)
|
||
from mdleaves import read_markdown, to_all, walk_markdown # noqa: E402
|
||
from mailleafs import from_mail_root # noqa: E402
|
||
|
||
CORPUS_DEFAULTS = ["README.md", "PLAN.md", "AGENTS.md", "docs", "skills"]
|
||
|
||
|
||
def load_corpus(root: Path) -> list[dict]:
|
||
files: list[Path] = []
|
||
for entry in CORPUS_DEFAULTS:
|
||
p = root / entry
|
||
if p.is_file():
|
||
files.append(p)
|
||
elif p.is_dir():
|
||
files.extend(walk_markdown(p))
|
||
leafs: list[dict] = []
|
||
for path in files:
|
||
try:
|
||
leafs.extend(to_all(read_markdown(path), path, repo="eSlider/2dph"))
|
||
except OSError as e:
|
||
print(f"kb/index: skip {path}: {e}", file=sys.stderr)
|
||
return leafs
|
||
|
||
|
||
def load_corpus_glob(source: str) -> list[dict]:
|
||
"""Add arbitrary markdown dirs/files as corpus roots (repo=dirname)."""
|
||
root = Path(source)
|
||
if not root.exists():
|
||
print(f"kb/index: skip missing corpus {source}", file=sys.stderr)
|
||
return []
|
||
files = [root] if root.is_file() else walk_markdown(root)
|
||
repo = root.name if root.is_dir() else root.parent.name
|
||
leafs: list[dict] = []
|
||
for path in files:
|
||
try:
|
||
leafs.extend(to_all(read_markdown(path), path, repo=repo))
|
||
except OSError as e:
|
||
print(f"kb/index: skip {path}: {e}", file=sys.stderr)
|
||
# yaml seeds (knowledge-mesh, workspace catalogs) as plain info leafs
|
||
if root.is_dir():
|
||
for path in sorted(root.rglob("*.y*ml")):
|
||
try:
|
||
leafs.append({
|
||
"source": str(path), "repo": repo, "heading": path.stem,
|
||
"text": path.read_text(encoding="utf-8", errors="replace")[:20000],
|
||
"type": "seed", "status": "current", "related": "",
|
||
})
|
||
except OSError:
|
||
continue
|
||
return leafs
|
||
|
||
|
||
def index_leafs(conn, leafs: list[dict], embed_fn, limit: int) -> tuple[int, int]:
|
||
count = 0
|
||
for lf in leafs[:limit] if limit else leafs:
|
||
query = f"{lf['heading']}\n\n{lf['text']}"
|
||
emb = embed_fn(lf["text"]) if lf["text"] else None
|
||
lid = upsert_leaf(conn, text=query, root="info", confidence="confirmed",
|
||
source=lf["source"], source_rev="working-tree",
|
||
how="kb/index", loc=lf["source"], type_=lf.get("type", "reference"),
|
||
embedding=emb)
|
||
link_from_file(conn, lid, lf["source"], repo=str(lf.get("repo") or ""))
|
||
count += 1
|
||
return count, len(leafs)
|
||
|
||
|
||
def embedder():
|
||
from model2vec import StaticModel
|
||
model = StaticModel.from_pretrained("minishlab/potion-multilingual-128M")
|
||
return lambda text: model.encode([text])[0].astype(float).tolist()
|
||
|
||
|
||
def index_fact_dicts(conn, facts: list[dict], embed_fn) -> int:
|
||
"""Write extract-shaped dicts as root=facts leafs (2-source source field)."""
|
||
leafs = []
|
||
for f in facts:
|
||
text = str(f.get("text") or "")
|
||
source = str(f.get("source") or "")
|
||
if not text or not source:
|
||
continue
|
||
leafs.append({
|
||
"text": text,
|
||
"root": "facts",
|
||
"confidence": "confirmed",
|
||
"source": source,
|
||
"source_rev": f.get("source_rev") or "working-tree",
|
||
"how": f.get("how") or "facts/extract",
|
||
"loc": f.get("loc") or source,
|
||
"type": "fact",
|
||
"embedding": embed_fn(text) if text else None,
|
||
})
|
||
return len(add_leafs(conn, leafs))
|
||
|
||
|
||
def facts_from_extract() -> list[dict]:
|
||
import subprocess
|
||
proc = subprocess.run(
|
||
[sys.executable, str(ROOT / "bin" / "facts" / "extract"), "--json", "--dry-run"],
|
||
cwd=ROOT,
|
||
capture_output=True,
|
||
text=True,
|
||
check=False,
|
||
)
|
||
if proc.returncode != 0:
|
||
print(f"kb/index: facts/extract failed: {proc.stderr}", file=sys.stderr)
|
||
return []
|
||
try:
|
||
payload = json.loads(proc.stdout)
|
||
except json.JSONDecodeError:
|
||
print("kb/index: facts/extract produced non-JSON", file=sys.stderr)
|
||
return []
|
||
return list(payload.get("facts") or [])
|
||
|
||
|
||
def main(argv: list[str]) -> int:
|
||
import argparse
|
||
p = argparse.ArgumentParser(description="build the 2dph brain index")
|
||
p.add_argument("--corpus", action="append", help="extra markdown dir/file to index (may repeat)")
|
||
p.add_argument("--rebuild", action="store_true", help="fresh db + indexes")
|
||
p.add_argument("--db", default="", help="path to kb.lbug (default var/kb.lbug)")
|
||
p.add_argument("--no-defaults", action="store_true", help="do not index repo README/docs/skills")
|
||
p.add_argument("--with-mail", action="store_true", help="include var/mail message.md leafs")
|
||
p.add_argument("--with-facts", action="store_true", help="run facts/extract into root=facts")
|
||
p.add_argument("--facts-json", default="", help="JSON list (or {facts:[...]}) of fact dicts")
|
||
p.add_argument(
|
||
"--with-chats",
|
||
nargs="?",
|
||
const=str(ROOT / "var" / "chats" / "md"),
|
||
default="",
|
||
help="index chat markdown as info (default var/chats/md)",
|
||
)
|
||
p.add_argument("--since", default="", help="with --with-mail, only messages dated >= YYYY-MM-DD")
|
||
p.add_argument("--dry-run", action="store_true", help="count leafs, write nothing")
|
||
p.add_argument(
|
||
"--skip-indexes",
|
||
action="store_true",
|
||
help="write leafs only; caller runs ensure_indexes after seeding facts",
|
||
)
|
||
p.add_argument("--limit", type=int, default=0, help="max leafs to embed")
|
||
p.add_argument("--json", action="store_true")
|
||
a = p.parse_args(argv)
|
||
|
||
from kblib import DB_PATH, VAR
|
||
|
||
dbpath = Path(a.db) if a.db else DB_PATH
|
||
leafs: list[dict] = [] if a.no_defaults else load_corpus(ROOT)
|
||
if a.corpus:
|
||
for source in a.corpus:
|
||
leafs.extend(load_corpus_glob(source))
|
||
chat_n = 0
|
||
if a.with_chats:
|
||
chats = load_corpus_glob(a.with_chats)
|
||
chat_n = len(chats)
|
||
leafs.extend(chats)
|
||
mail_n = 0
|
||
if a.with_mail:
|
||
mail = from_mail_root(ROOT / "var" / "mail", since=a.since)
|
||
mail_n = len(mail)
|
||
leafs.extend(mail)
|
||
|
||
facts: list[dict] = []
|
||
if a.facts_json:
|
||
raw = Path(a.facts_json).read_text(encoding="utf-8")
|
||
payload = json.loads(raw)
|
||
facts = list(payload.get("facts") if isinstance(payload, dict) else payload)
|
||
if a.with_facts:
|
||
facts.extend(facts_from_extract())
|
||
|
||
if a.dry_run:
|
||
msg = {
|
||
"indexed": 0,
|
||
"corpus_total": len(leafs),
|
||
"mail_leafs": mail_n,
|
||
"chat_leafs": chat_n,
|
||
"facts_leafs": len(facts),
|
||
"dry_run": True,
|
||
}
|
||
print(json.dumps(msg, indent=2) if a.json else
|
||
f"brain/index: {len(leafs)} info + {len(facts)} facts would be indexed")
|
||
return 0
|
||
|
||
VAR.mkdir(exist_ok=True)
|
||
dbpath.parent.mkdir(parents=True, exist_ok=True)
|
||
if a.rebuild and dbpath.exists():
|
||
dbpath.unlink()
|
||
|
||
db, conn = connect(dbpath, read_only=False)
|
||
init_schema(conn)
|
||
|
||
embed = embedder()
|
||
done, total = index_leafs(conn, leafs, embed, a.limit)
|
||
fact_n = index_fact_dicts(conn, facts, embed) if facts else 0
|
||
if not a.skip_indexes:
|
||
ensure_indexes(conn)
|
||
s = stats(conn)
|
||
conn.close()
|
||
db.close()
|
||
|
||
result = {
|
||
"indexed": done,
|
||
"corpus_total": total,
|
||
"facts_leafs": fact_n,
|
||
"chat_leafs": chat_n,
|
||
**{k: v for k, v in s.items() if k in ("total", "by_root")},
|
||
}
|
||
if a.skip_indexes:
|
||
result["indexes"] = "skipped"
|
||
print(json.dumps(result, indent=2) if a.json else
|
||
f"indexed {done}/{total} info + {fact_n} facts; db total {s['total']}")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main(sys.argv[1:])) |