cortexm/experimental/graph_recall.py
Rama main · réplica de lectura del repositorio original.
"""Entity→fact adjacency index with bounded graph walks (μ=0).
Borrowed from RuVector's graph-memory reconstruction ("follow
Cue → Tag → Content instead of one embedding search"), stripped to a
deterministic core:
* an inverted index maps entity tokens (from fact subjects and
values) to the fact ids they participate in;
* a query's tokens that HIT the index anchor a bounded BFS walk
(default 2 hops) that pulls facts CONNECTED to the query's
entities even when lexical/VSA similarity to the fact's triple is
low — the multi-session failure mode where a similar-looking fact
from the WRONG session wins on embedding similarity alone;
* facts reached through MULTIPLE distinct query tokens score higher
(cross-confirmation), and each hop decays the contribution.
Why this is Pareto for cortexm: the retrieval misses it targets do
not need semantic similarity at all — they need the join structure
the extractor already produced. The index is one pass over the facts
table; the walk is set intersection. No LLM, no learned weights,
byte-deterministic output.
"""
from __future__ import annotations
import threading
from typing import Any
# Minimal stopword set — kept local so the module has zero deps on
# the bench/judge layer (which lives under scripts/).
_STOPWORDS = frozenset("""
a an and are as at be been by did do does for from had has have how i
in is it its me my not of on or our so that the their them they this
to was we were what when where which who why will with you your
just got get very really much many name named called
""".split())
_MIN_TOKEN_LEN = 3
def _entity_tokens(subject: str, value: str) -> list[str]:
"""Deterministic entity tokens for one fact triple.
Subject: the raw subject plus its ``user:``-stripped form.
Value: whitespace/punctuation-split tokens with stopwords and
short tokens dropped (so ``"dog named Charlie"`` yields
``dog, charlie``).
"""
toks: list[str] = []
subj = (subject or "").strip()
if subj:
toks.append(subj.lower())
if ":" in subj:
toks.append(subj.split(":", 1)[1].lower())
val = (value or "").strip()
if val:
for piece in val.replace(",", " ").replace(".", " ").split():
piece = piece.strip("'\"!?;:()[]").lower()
if (len(piece) >= _MIN_TOKEN_LEN
and piece not in _STOPWORDS
and piece.isalnum()):
toks.append(piece)
return toks
class EntityGraphIndex:
"""Inverted entity-token → fact-id adjacency index.
Built lazily from the store's ACTIVE, non-quarantined facts for a
user; cached and rebuilt automatically when the underlying active
fact count changes (cheap generation check — one indexed COUNT
query). Thread-safe via an RLock around build/lookup.
"""
def __init__(self, store: Any) -> None:
self._store = store
self._lock = threading.RLock()
self._index: dict[str, set[str]] = {}
self._fact_entities: dict[str, list[str]] = {}
self._built_for: tuple[str | None, int] | None = None
# ------------------------------------------------------------------ build
def _ensure(self, user_id: str | None) -> None:
n = self._store.count_facts(user_id=user_id, active_only=True)
key = (user_id, n)
with self._lock:
if self._built_for == key and self._index:
return
facts = self._store.query_facts(user_id=user_id, active=True,
include_quarantined=False)
index: dict[str, set[str]] = {}
fact_entities: dict[str, list[str]] = {}
for f in facts:
ents = _entity_tokens(f.subject, f.value)
fact_entities[f.id] = ents
for e in ents:
index.setdefault(e, set()).add(f.id)
self._index = index
self._fact_entities = fact_entities
self._built_for = key
# ------------------------------------------------------------------ query
def query_tokens(self, query: str) -> list[str]:
"""Query tokens that anchor the walk (must exist in the index)."""
toks: list[str] = []
for piece in (query or "").replace(",", " ").replace("?", " ") \
.replace(".", " ").split():
piece = piece.strip("'\"!?;:()[]").lower()
if (len(piece) >= _MIN_TOKEN_LEN
and piece not in _STOPWORDS
and piece in self._index
and piece not in toks):
toks.append(piece)
return toks
def walk(self, query: str, user_id: str | None,
max_hops: int = 2, limit: int = 64,
scope: frozenset[str] | set[str] | None = None
) -> tuple[dict[str, float], dict]:
"""Bounded BFS from the query's entity tokens.
Returns (fact_id → raw score, stats). Scores:
hop 1: +1.0 per DISTINCT query token indexing the fact
hop 2: +0.5 per distinct hop-1 entity indexing the fact
Scores are NOT normalized here — the caller scales by its own
config weight. Facts outside ``scope`` (when given) are
dropped.
"""
self._ensure(user_id)
stats = {"n_query_tokens": 0, "hop1": 0, "hop2": 0}
with self._lock:
anchors = self.query_tokens(query)
stats["n_query_tokens"] = len(anchors)
if not anchors:
return {}, stats
scores: dict[str, float] = {}
# ---- hop 1: facts directly indexed by a query token ----
hop1_entities: set[str] = set()
for tok in anchors:
for fid in self._index.get(tok, ()):
if scope is not None and fid not in scope:
continue
scores[fid] = scores.get(fid, 0.0) + 1.0
stats["hop1"] += 1
hop1_entities.update(self._fact_entities.get(fid, ()))
# ---- hop 2: facts of hop-1 entities (anchors excluded) ----
if max_hops >= 2:
n_hop2 = 0
for ent in sorted(hop1_entities): # sorted → deterministic
if ent in anchors:
continue
for fid in sorted(self._index.get(ent, ())):
if scope is not None and fid not in scope:
continue
if fid in scores:
continue # hop-1 facts keep their higher score
scores[fid] = 0.5
n_hop2 += 1
if n_hop2 >= limit:
break
if n_hop2 >= limit:
break
stats["hop2"] = n_hop2
return scores, stats
# ---------------------------------------------------------------------------
# Reader-facing facade (wired in bridge/reader.py behind
# Config.graph_recall_enabled)
# ---------------------------------------------------------------------------
_ACTIVE: dict[int, EntityGraphIndex] = {}
def graph_recall_boost(store: Any, query: str, user_id: str | None,
scope: frozenset[str] | set[str] | None = None,
max_hops: int = 2
) -> tuple[dict[str, float], dict]:
"""Score boost map for one query, μ=0.
Uses one process-wide index per store instance (id() keyed — the
reader holds its store for its lifetime; the generation check
inside rebuilds when facts change).
"""
key = id(store)
idx = _ACTIVE.get(key)
if idx is None:
idx = EntityGraphIndex(store)
_ACTIVE[key] = idx
# opportunistic cleanup: drop dead-store indexes
if len(_ACTIVE) > 8:
for k in list(_ACTIVE.keys()):
if k != key and _ACTIVE[k]._built_for is None:
_ACTIVE.pop(k, None)
return idx.walk(query, user_id, max_hops=max_hops, scope=scope)