echo/echo-memory.plugin.src/skills/echo-memory/scripts/echo_recall.py

#!/usr/bin/env python3
"""echo_recall.py — H1: hybrid lexical (BM25) + graph recall.  [v1.0 — wired]

Replaces the keyword-only recall (search/simple + fixed one-hop) with a ranked, local
retriever that finds a note even when the query is phrased differently than the note
was stored, then expands along the graph with a per-hop decay so the *connected*
context is surfaced in relevance order.

DESIGN (bounded by the project's pure-stdlib, no-deps constraint):
  * Lexical layer  : BM25 over note bodies. No embeddings / external API — BM25 is a
                     strong, dependency-free baseline that handles vocabulary mismatch
                     far better than substring search.
  * Graph layer    : from the top lexical hits, walk body wikilinks + source_notes up
                     to MAX_HOPS, scoring each node `parent * GRAPH_DECAY**hop` (max over
                     paths) so near, strongly-cued neighbours rank above distant ones.
  * Persistence    : the BM25 postings + doc stats are a DERIVED artifact, so they live
                     in the vault's machine index dir (`_agent/index/recall-index.json`),
                     maintained on `capture` (update_note) and fully rebuilt by `sweep.py`
                     — exactly how `entities.json` is maintained. Rebuildable => portable.
  * Resilience     : if the index can't be built (vault unreachable), recall falls back
                     to the server's /search/simple so it degrades instead of dying.

Corpus = the indexable ENTITY notes (idx_mod.kind_for_path is not None) — people,
companies, concepts, references, meetings, projects, areas, decisions, semantic/
episodic memory, skills — PLUS the time-series notes where conversation-derived
context actually lives: session logs (`_agent/sessions/`) and journal notes
(`journal/daily|weekly|monthly|quarterly|annual/`). Time-series docs carry a
per-kind DOWN-WEIGHT so entity notes still rank first for the same lexical match.

RANKING (v1.5): final score = BM25 * kind_weight * freshness * status_factor —
  * kind_weight : 1.0 entity · 0.5 session · 0.4 journal (KIND_WEIGHT)
  * freshness   : half-life decay on `updated:` (or the filename date), floored at
                  FRESH_FLOOR so old notes are demoted, never buried
  * status      : `active` boosted, `archived` demoted (STATUS_FACTOR)
so a stale archived project no longer outranks the live one on raw term frequency.
Doc metadata rides in the index (schema 2); a schema-1 index is discarded and
rebuilt on the next recall/sweep.
"""
from __future__ import annotations

import datetime as _dt
import json
import math
import os
import re
import sys
import urllib.parse
from collections import Counter, deque
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import echo                      # noqa: E402
import echo_index as idx_mod     # noqa: E402
import echo_links as links       # noqa: E402

RECALL_INDEX_PATH = "_agent/index/recall-index.json"
INDEX_SCHEMA = 2                 # bumped: schema 2 adds per-doc meta (weight/updated/status)

# Tuning knobs. Defaults are standard BM25 (k1/b) + a 0.6 hop decay over <=2 hops.
K1 = 1.5
B = 0.75
MAX_HOPS = 2
GRAPH_DECAY = 0.6
MIN_SCORE = 0.0  # relevance floor; raise to trim weak lexical hits

# --- rank-fusion priors -------------------------------------------------------
# Time-series notes are in the corpus but down-weighted: they are where "what did we
# decide three weeks ago" lives, yet the entity note should win a tie on the same terms.
_TIME_SERIES = [
    (re.compile(r"^_agent/sessions/"), "session"),
    (re.compile(r"^journal/daily/"), "daily"),
    (re.compile(r"^journal/(weekly|monthly|quarterly|annual)/"), "rollup"),
]
KIND_WEIGHT = {"session": 0.5, "daily": 0.4, "rollup": 0.4}
STATUS_FACTOR = {"active": 1.1, "on-hold": 0.9, "archived": 0.75}
FRESH_HALF_LIFE = float(os.environ.get("ECHO_FRESH_HALF_LIFE", "90"))  # days
FRESH_FLOOR = 0.6   # freshness multiplier never drops below this — demote, don't bury

_FM_UPDATED = re.compile(r"(?m)^updated:\s*[\"']?(\d{4}-\d{2}-\d{2})")
_FM_STATUS = re.compile(r"(?m)^status:\s*[\"']?([A-Za-z-]+)")
_PATH_DATE = re.compile(r"(\d{4}-\d{2}-\d{2})")


def _series_kind(path: str) -> str | None:
    for rx, kind in _TIME_SERIES:
        if rx.match(path):
            return kind
    return None


def doc_meta(path: str, raw_text: str) -> dict:
    """Per-doc ranking priors, extracted once at index time: kind weight, last-updated
    date (frontmatter `updated:`, else the date in the filename), and `status:`."""
    kind = _series_kind(path)
    w = KIND_WEIGHT.get(kind, 1.0) if kind else 1.0
    head = raw_text[:2000]
    mu = _FM_UPDATED.search(head)
    updated = mu.group(1) if mu else None
    if updated is None:
        mp = _PATH_DATE.search(path.rsplit("/", 1)[-1])
        updated = mp.group(1) if mp else None
    ms = _FM_STATUS.search(head)
    meta = {"w": w}
    if updated:
        meta["u"] = updated
    if ms:
        meta["s"] = ms.group(1)
    return meta


def freshness(updated: str | None, today_s: str) -> float:
    """Half-life decay on document age, floored. Unknown age is neutral (1.0)."""
    if not updated:
        return 1.0
    try:
        age = (_dt.date.fromisoformat(today_s) - _dt.date.fromisoformat(updated)).days
    except ValueError:
        return 1.0
    if age <= 0:
        return 1.0
    return FRESH_FLOOR + (1.0 - FRESH_FLOOR) * (0.5 ** (age / FRESH_HALF_LIFE))

_WORD = re.compile(r"[a-z0-9]+")
# Minimal English stoplist — keeps BM25 idf from being dominated by glue words.
_STOP = frozenset(
    "the a an and or of to in is are was were be been for on at by with as it this that "
    "from into over under not no yes do does did has have had will would can could".split()
)
_TEMPLATE_RE = re.compile(r"(^|/)(templates/|.*-template\.md$)")
_SKIP_BASENAMES = {"README.md"}


def tokenize(text: str) -> list[str]:
    """Lowercase word tokens, stopworded. Frontmatter should be stripped by the caller."""
    return [t for t in _WORD.findall(text.lower()) if t not in _STOP and len(t) > 1]


def strip_frontmatter(text: str) -> str:
    if not text.startswith("---"):
        return text
    end = text.find("\n---", 3)
    return text[end + 4:] if end != -1 else text


# --------------------------------------------------------------------- BM25 core
class Bm25Index:
    """A compact, serializable BM25 index with per-doc ranking meta.
    add()/remove() are idempotent per path. add() takes the RAW note text — it strips
    frontmatter itself and extracts the doc meta (weight/updated/status) in one pass."""

    def __init__(self) -> None:
        self.df: Counter[str] = Counter()              # term -> #docs containing it
        self.postings: dict[str, dict[str, int]] = {}  # term -> {doc_path: tf}
        self.length: dict[str, int] = {}               # doc_path -> token count
        self.meta: dict[str, dict] = {}                # doc_path -> {w, u, s} priors
        self.n_docs = 0
        self.avg_len = 0.0

    def _recompute(self) -> None:
        self.n_docs = len(self.length)
        self.avg_len = (sum(self.length.values()) / self.n_docs) if self.n_docs else 0.0

    def add(self, path: str, raw_text: str) -> None:
        if path in self.length:        # re-index in place (idempotent)
            self.remove(path)
        toks = tokenize(strip_frontmatter(raw_text))
        self.length[path] = len(toks)
        self.meta[path] = doc_meta(path, raw_text)
        for term, tf in Counter(toks).items():
            self.postings.setdefault(term, {})[path] = tf
            self.df[term] += 1
        self._recompute()

    def remove(self, path: str) -> None:
        if path not in self.length:
            return
        for term in list(self.postings.keys()):
            if path in self.postings[term]:
                del self.postings[term][path]
                self.df[term] -= 1
                if not self.postings[term]:
                    del self.postings[term]
                    del self.df[term]
        del self.length[path]
        self.meta.pop(path, None)
        self._recompute()

    def _doc_factor(self, path: str, today_s: str) -> float:
        """The rank-fusion prior: kind weight x freshness x status factor."""
        m = self.meta.get(path) or {}
        return (float(m.get("w", 1.0))
                * freshness(m.get("u"), today_s)
                * STATUS_FACTOR.get(m.get("s", ""), 1.0))

    def score(self, query: str, limit: int = 10,
              today_s: str | None = None) -> list[tuple[str, float]]:
        today_s = today_s or echo.today()
        scores: Counter[str] = Counter()
        for term in set(tokenize(query)):
            posting = self.postings.get(term)
            if not posting:
                continue
            idf = math.log(1 + (self.n_docs - self.df[term] + 0.5) / (self.df[term] + 0.5))
            for path, tf in posting.items():
                dl = self.length.get(path, 0)
                denom = tf + K1 * (1 - B + B * (dl / self.avg_len if self.avg_len else 1))
                scores[path] += idf * (tf * (K1 + 1) / denom if denom else 0)
        fused = {p: s * self._doc_factor(p, today_s) for p, s in scores.items()}
        ranked = sorted(fused.items(), key=lambda kv: -kv[1])[:limit]
        return [(p, s) for p, s in ranked if s > MIN_SCORE]

    # ---- serialization (compact; rebuildable, so the format can change freely) ----
    def to_json(self) -> dict:
        return {"schema": INDEX_SCHEMA, "n_docs": self.n_docs, "avg_len": self.avg_len,
                "length": self.length, "postings": self.postings, "meta": self.meta}

    @classmethod
    def from_json(cls, d: dict) -> "Bm25Index":
        if d.get("schema") != INDEX_SCHEMA:
            return cls()  # older/newer format -> empty; recall's cold path rebuilds
        ix = cls()
        ix.length = d.get("length", {})
        ix.postings = d.get("postings", {})
        ix.meta = d.get("meta", {})
        ix.n_docs = d.get("n_docs", len(ix.length))
        ix.avg_len = d.get("avg_len", 0.0)
        ix.df = Counter({term: len(posting) for term, posting in ix.postings.items()})
        return ix


# ------------------------------------------------------------------- vault I/O
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"recall 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 _indexable(path: str) -> bool:
    base = path.rsplit("/", 1)[-1]
    if (not path.endswith(".md") or base in _SKIP_BASENAMES
            or _TEMPLATE_RE.search(path)):
        return False
    # Entity notes (the durable graph) + time-series notes (sessions, journal) — the
    # latter join the corpus down-weighted so past decisions/discussions are findable.
    return idx_mod.kind_for_path(path) is not None or _series_kind(path) is not None


# ------------------------------------------------------------------- persistence
def load_index() -> Bm25Index:
    status, body = echo.request("GET", echo.vault_url(RECALL_INDEX_PATH))
    if status == 404:
        return Bm25Index()
    echo.check(status, body, "recall-index load")
    try:
        return Bm25Index.from_json(json.loads(body))
    except (json.JSONDecodeError, KeyError, TypeError):
        return Bm25Index()


def save_index(ix: Bm25Index) -> None:
    body = json.dumps(ix.to_json(), ensure_ascii=False).encode("utf-8")
    status, b = echo.request("PUT", echo.vault_url(RECALL_INDEX_PATH), data=body,
                             headers={"Content-Type": "application/json"})
    echo.check(status, b, "recall-index save")


def rebuild(prefetched: dict | None = None) -> Bm25Index:
    """Full rebuild from every indexable entity note. Called by sweep.py and as the
    cold-cache path inside recall(). Saves and returns the index.

    `prefetched` (a {path: text} map, e.g. from echo.read_many) lets sweep.py reuse the
    content it already pulled instead of walking and re-fetching the whole vault again."""
    ix = Bm25Index()
    if prefetched is not None:
        items = ((p, t) for p, t in prefetched.items() if _indexable(p))
    else:
        items = ((p, _get(p)) for p in _walk() if _indexable(p))
    for path, text in items:
        if text is not None:
            ix.add(path, text)   # add() strips frontmatter + extracts doc meta itself
    save_index(ix)
    return ix


def update_note(path: str, text: str) -> None:
    """Best-effort incremental upkeep for a single note (mirrors how entities.json is
    maintained on capture). Never raises into the caller. The load->save is wrapped in the
    advisory lock with a fresh re-read (H3), so concurrent captures can't clobber the
    recall index either."""
    try:
        if not _indexable(path):
            return  # only entity notes are part of the recall corpus
        import echo_concurrency
        with echo_concurrency.vault_lock():
            ix = load_index()                      # fresh read inside the lock
            ix.add(path, text)                     # raw text: add() strips + extracts meta
            save_index(ix)
    except Exception as exc:  # noqa: BLE001 — upkeep must never fail a capture
        print(f"echo_recall: index update skipped ({exc})", file=sys.stderr)


# ------------------------------------------------------------------- graph fusion
def _resolve_target(target: str, nmap: dict) -> str | None:
    target = target.strip()
    if not target:
        return None
    if "/" in target:
        return target if target.endswith(".md") else target + ".md"
    return nmap.get(idx_mod.slugify(target))


def expand_graph(seeds, nmap: dict, base_scores: dict, max_hops: int = MAX_HOPS):
    """BFS from seeds over body wikilinks + source_notes, scoring each neighbour
    `parent_score * GRAPH_DECAY**hop` (kept as the max over reaching paths). Returns a
    score-sorted list of (path, (score, via)) excluding the seeds themselves."""
    score_of = dict(base_scores)
    seen = set(seeds)
    results: dict[str, tuple[float, str]] = {}
    # Breadth-first by HOP so each frontier's note bodies can be fetched concurrently
    # (echo.read_many) instead of one blocking GET per node. Scoring is unchanged: a
    # neighbour keeps the max decayed score over the paths that reach it, and non-seed
    # parents score 1.0 exactly as the serial deque version did.
    frontier = list(seeds)
    for hop in range(max_hops):
        if not frontier:
            break
        texts = echo.read_many(frontier)
        next_frontier: list[str] = []
        for path in frontier:
            text = texts.get(path)
            if text is None:
                continue
            body = strip_frontmatter(text)
            targets = set(links.all_wikilinks(body)) | set(links.source_notes(text))
            parent = score_of.get(path, 1.0)
            decayed = parent * (GRAPH_DECAY ** (hop + 1))
            for t in targets:
                tp = _resolve_target(t, nmap)
                if not tp or tp in seeds:
                    continue
                if decayed > results.get(tp, (0.0, ""))[0]:
                    results[tp] = (decayed, path)
                if tp not in seen:
                    seen.add(tp)
                    next_frontier.append(tp)
        frontier = next_frontier
    return sorted(results.items(), key=lambda kv: -kv[1][0])


# ------------------------------------------------------------------- presentation
def _doc_summary(path: str, text: str) -> dict:
    """Structured per-note summary shared by the human brief and --json output."""
    out: dict = {"path": path}
    mtype = re.search(r"(?m)^type:\s*(.+)$", text)
    if mtype:
        out["type"] = mtype.group(1).strip()
    mu = _FM_UPDATED.search(text[:2000])
    if mu:
        out["updated"] = mu.group(1)
    ms = _FM_STATUS.search(text[:2000])
    if ms:
        out["status"] = ms.group(1)
    body = strip_frontmatter(text)
    status = re.search(r"(?ms)^##\s+Status\s*$(.*?)(?=^##\s|\Z)", body)
    if status and status.group(1).strip():
        out["excerpt"] = status.group(1).strip()[:400]
    else:
        kept = [ln[:200] for ln in body.splitlines()
                if ln.strip() and not ln.startswith("# ")][:6]
        out["excerpt"] = "\n".join(kept)
    return out


def _brief(path: str, score: float | None = None, via: str | None = None) -> None:
    text = links.get_text(path)
    if text is None:
        return
    info = _doc_summary(path, text)
    head = f"\n### {path}"
    if score is not None:
        head += f"   (score {score:.2f})"
    if via:
        head += f"   (via {via})"
    print(head)
    stamps = []
    if info.get("type"):
        stamps.append(f"type: {info['type']}")
    if info.get("updated"):
        stamps.append(f"updated: {info['updated']}")
    if info.get("status"):
        stamps.append(f"status: {info['status']}")
    if stamps:
        print("_" + " · ".join(stamps) + "_")   # staleness is visible, not guessed
    if info.get("excerpt"):
        print(info["excerpt"])


# ------------------------------------------------------------------- entrypoint
def recall(query, limit: int = 8, as_json: bool = False) -> int:
    q = " ".join(query) if isinstance(query, list) else query
    today_s = echo.today()
    index = idx_mod.load()
    nmap = idx_mod.name_map(index)

    # --- lexical layer (BM25); rebuild the cache on a cold miss --------------
    lexical: list[tuple[str, float]] = []
    try:
        ix = load_index()
        if ix.n_docs == 0:
            ix = rebuild()
        lexical = ix.score(q, limit=limit, today_s=today_s)
    except echo.EchoError as exc:
        print(f"recall: BM25 index unavailable ({exc}); falling back to server search",
              file=sys.stderr)

    base: dict[str, float] = {p: s for p, s in lexical}
    hits: list[str] = [p for p, _ in lexical]

    # --- entity-index seed (alias-aware exact match) -------------------------
    _, e = idx_mod.resolve(index, q)
    if e and e.get("path") and e["path"] not in base:
        hits.insert(0, e["path"])
        base[e["path"]] = max(base.values(), default=1.0)

    # --- resilience fallback: server search when lexical found nothing -------
    if not hits:
        try:
            status, body = echo.request(
                "POST", f"{echo.BASE}/search/simple/?query={urllib.parse.quote_plus(q)}")
            if status == 200:
                for r in (json.loads(body) or [])[:limit]:
                    fn = r.get("filename") or r.get("path")
                    if fn and fn not in base:
                        hits.append(fn)
                        base[fn] = 0.0
        except Exception:  # noqa: BLE001 — fallback only; stay quiet
            pass

    # --- graph layer ----------------------------------------------------------
    neighbours = expand_graph(hits, nmap, base, max_hops=MAX_HOPS)

    if as_json:
        import echo_output
        texts = echo.read_many(hits + [p for p, _ in neighbours[: 2 * limit]])
        primary = []
        for p in hits:
            if texts.get(p) is None:
                continue
            primary.append({**_doc_summary(p, texts[p]),
                            "score": round(base.get(p, 0.0), 3)})
        linked = []
        for p, (sc, via) in neighbours[: 2 * limit]:
            if texts.get(p) is None:
                continue
            linked.append({**_doc_summary(p, texts[p]),
                           "score": round(sc, 3), "via": via})
        env = echo_output.envelope("recall", {"query": q, "primary": primary,
                                              "linked": linked})
        print(json.dumps(env, ensure_ascii=False))
        return 0

    print(f"== recall: {q} ==")
    print(f"\n# Primary hits ({len(hits)})")
    for p in hits:
        _brief(p, score=base.get(p))
    print(f"\n# Linked context ({len(neighbours)})")
    for p, (sc, via) in neighbours[: 2 * limit]:
        _brief(p, score=sc, via=via)
    return 0