chorus/echo-memory.plugin.src/skills/echo-memory/scripts/sweep.py

#!/usr/bin/env python3
"""sweep.py — bring an existing ECHO vault up to the current feature spec.

After upgrading the plugin (entity index, cross-linking, recall), run this once to
backfill what older vaults don't have yet:

  1. ensure the `_agent/index/` folder exists,
  2. **(re)build the entity index** (`_agent/index/entities.json`) from the notes
     already in the vault (people, companies, projects, concepts, references,
     meetings, areas, decisions, semantic/episodic memory, skills),
  3. **symmetrize cross-links** — for every `## Related` link A -> B, ensure the
     reciprocal B -> A exists (it only adds the missing direction; it never invents
     new links from body text), and
  4. stamp the marker `schema_version` to the current schema (3).

Dry-run by default; pass --apply to write. READ-ONLY without --apply.
Cross-platform: pure Python via echo.py.

Usage:  sweep.py [--apply]
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR))

import echo          # noqa: E402
import echo_index as idx_mod   # noqa: E402
import echo_links as links     # noqa: E402
import echo_recall             # noqa: E402

CURRENT_SCHEMA = 4
TEMPLATE_RE = re.compile(r"(^|/)(templates/|.*-template\.md$)")
SKIP_BASENAMES = {"README.md"}


kind_for = idx_mod.kind_for_path


def get(path: str) -> str | None:
    status, body = echo.request("GET", echo.vault_url(path))
    if status == 404:
        return None
    echo.check(status, body, f"get {path}")
    return body.decode(errors="replace")


def list_dir(path: str):
    p = "" if path in ("", "/") else (path if path.endswith("/") else path + "/")
    body = get(p)
    if body is None:
        return [], []
    try:
        j = json.loads(body)
    except json.JSONDecodeError:
        return [], []
    entries = list(j.get("files", [])) + list(j.get("folders", []))
    files = [e for e in entries if not e.endswith("/")]
    folders = [e[:-1] for e in entries if e.endswith("/")]
    return files, folders


def walk(prefix: str = ""):
    files, folders = list_dir(prefix)
    for f in files:
        yield prefix + f
    for d in folders:
        yield from walk(f"{prefix}{d}/")


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Bring an ECHO vault up to the current feature spec")
    ap.add_argument("--apply", action="store_true")
    args = ap.parse_args(argv)
    apply = args.apply
    tag = "APPLY" if apply else "PLAN "

    if get("_agent/echo-vault.md") is None:
        print("sweep: marker missing — vault not bootstrapped. Run bootstrap.py first.", file=sys.stderr)
        return 3

    print(f"sweep: {'applying changes' if apply else 'dry-run (no writes)'}\n")

    # ---- 1. index folder ------------------------------------------------------
    if get("_agent/index/README.md") is None:
        print(f"sweep: {tag} create _agent/index/README.md")
        if apply:
            echo.request("PUT", echo.vault_url("_agent/index/README.md"),
                         data=b"# index\n\nMachine-maintained entity index. See the echo-memory plugin.\n",
                         headers={"Content-Type": "text/markdown"})

    all_files = [p for p in walk()
                 if p.endswith(".md") and p.rsplit("/", 1)[-1] not in SKIP_BASENAMES
                 and not TEMPLATE_RE.search(p)]

    # Fetch every note's content ONCE, concurrently. All three passes below (index,
    # recall, link symmetrize) read from this shared cache, so no file is fetched
    # twice and link targets aren't re-fetched per reference.
    texts = echo.read_many(all_files)
    print(f"sweep: read {len(texts)} notes (concurrent, keep-alive)\n")

    # ---- 2. (re)build entity index -------------------------------------------
    index = idx_mod.load()
    existing = dict(index.get("entities", {}))
    rebuilt = idx_mod.empty_index()
    added = updated = 0
    for path in all_files:
        kind = kind_for(path)
        if kind is None:
            continue
        base = path.rsplit("/", 1)[-1][:-3]
        slug = idx_mod.slugify(base)
        text = texts.get(path) or ""
        title = links.first_h1(text) or base
        prev = existing.get(slug)
        # Frontmatter aliases are authoritative and re-folded every rebuild; prior index
        # aliases (e.g. mentions learned by capture) are preserved; upsert adds title variants.
        prev_aliases = prev.get("aliases", []) if prev else []
        aliases = list(prev_aliases) + idx_mod.aliases_in_frontmatter(text)
        idx_mod.upsert(rebuilt, slug, path, kind, title, aliases)
        if not prev:
            added += 1
        elif prev.get("path") != path or prev.get("title") != title:
            updated += 1
    print(f"sweep: {tag} entity index — {len(rebuilt['entities'])} entities "
          f"({added} new, {updated} changed)")
    if apply:
        idx_mod.save(rebuilt)

    # ---- 2b. (re)build the recall (BM25) index -------------------------------
    if apply:
        rix = echo_recall.rebuild(prefetched=texts)
        print(f"sweep: {tag} recall index — {rix.n_docs} docs (BM25)")
    else:
        n_idx = sum(1 for p in all_files if echo_recall._indexable(p))
        print(f"sweep: {tag} recall index — would rebuild BM25 over {n_idx} entity notes")

    # ---- 3. symmetrize existing Related links --------------------------------
    fileset = set(all_files)
    basemap: dict[str, str] = {}
    for p in all_files:
        basemap.setdefault(idx_mod.slugify(p.rsplit("/", 1)[-1][:-3]), p)

    def resolve_target(t: str) -> str | None:
        t = t.strip()
        if "/" in t:
            cand = t if t.endswith(".md") else t + ".md"
            return cand if cand in fileset else None
        return basemap.get(idx_mod.slugify(t))

    missing = set()  # (target_path, source_path) — reciprocal link to add on target
    for path in all_files:
        text = texts.get(path)
        if text is None:
            continue
        for tgt in links.related_targets(text):
            tp = resolve_target(tgt)
            if not tp or tp == path:
                continue
            back = {resolve_target(t) for t in links.related_targets(texts.get(tp) or "")}
            if path not in back:
                missing.add((tp, path))
    missing = sorted(missing)
    print(f"sweep: {tag} reciprocal links — {len(missing)} to add")
    for tp, sp in missing[:40]:
        print(f"          {tp}  +=  [[{links.link_token(sp)}]]")
    if len(missing) > 40:
        print(f"          ... and {len(missing) - 40} more")
    if apply:
        for tp, sp in missing:
            links.add_one(tp, sp)

    # ---- 4. stamp schema ------------------------------------------------------
    marker = get("_agent/echo-vault.md") or ""
    cur = next((int(m) for m in re.findall(r"(?m)^schema_version:\s*(\d+)", marker)), 0)
    if cur < CURRENT_SCHEMA:
        print(f"sweep: {tag} schema_version {cur} -> {CURRENT_SCHEMA}")
        if apply:
            echo.cmd_fm("_agent/echo-vault.md", "schema_version", str(CURRENT_SCHEMA))

    print(f"\nsweep: {'done — run vault_lint.py to confirm invariants' if apply else 'dry-run complete. Re-run with --apply to write.'}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())