class CortexMemory:
"""Persistent memory engine: memorize / recall / list / forget for one user."""
def __init__(
self,
provider: LLMProvider,
store: MemoryStore,
user_id: str = "local",
top_k: int = 5,
extractor: LLMProvider | None = None,
use_episodic: bool = False,
use_graph: bool = False,
use_dedup: bool = False,
dedup_threshold: float = 0.95,
use_soft_update: bool = False,
arbiter: LLMProvider | None = None,
reranker: Reranker | None = None,
rerank_pool: int = 60,
) -> None:
self.provider = provider
self.store: MemoryStore = store
self.user_id = user_id
self.top_k = top_k
# L4 episodic (additive, OFF by default): with ``use_episodic`` false or no ``extractor``,
# memorize/recall run the exact pre-episodic path — zero extra calls, byte-identical.
self.extractor = extractor
self.use_episodic = use_episodic
# G2 write-time entity graph (additive, OFF by default). With ``use_graph`` false (or no
# ``extractor``) memorize writes NO entity/edge/link rows and is byte-identical to today. It
# reuses the SAME folded extraction the episodic path already makes, so enabling both spends
# a single flash-lite call per memorize. Recall never reads these tables.
self.use_graph = use_graph
# L5 anti-saturation (additive, OFF by default). With BOTH ``use_dedup`` and
# ``use_soft_update`` false, memorize AND recall run the exact pre-L5 path — no neighbour
# scan, no arbiter call, and recall never queries ``superseded_ids`` (byte-identical). Dedup
# is embedding-only (no LLM); soft-update spends ONE cheap ``arbiter`` call only inside the
# "related but not duplicate" band and only when an ``arbiter`` is configured.
self.use_dedup = use_dedup
self.dedup_threshold = dedup_threshold
self.use_soft_update = use_soft_update
self.arbiter = arbiter
# Hosted cross-encoder reranker (additive, OFF by default). With ``reranker`` None, recall
# runs the EXACT pre-rerank path — same retrieval depth (``k``), no reorder, byte-identical.
# When set, recall over-retrieves a POOL of ``max(rerank_pool, k)`` candidates and asks the
# reranker to select the top ``k``; ANY reranker failure degrades to the RRF top-``k`` pool,
# so recall can never fail because of reranking.
self.reranker = reranker
self.rerank_pool = rerank_pool
# The embedding-signature guard is enforced LAZILY (on memorize/recall), NOT here, so a
# model/dim mismatch can never brick the whole server on startup — list_memories and
# forget stay usable to inspect and recover the store.
self._embed_model, self._embed_dim = _embed_signature(provider)
def _embed(self, text: str, what: str) -> list[float]:
"""Verify the embedding signature, embed one text, and reject an empty/wrong-dim result.
The signature guard runs here (not in __init__) so a model/dim mismatch only blocks
embed-dependent ops, leaving list/forget usable to recover. A partial provider failure
can return no vector; storing/querying an empty/short vector silently corrupts
retrieval, so we fail loudly instead.
"""
self.store.ensure_embedding(self._embed_model, self._embed_dim)
vectors = self.provider.embed([text]).vectors
vec = list(vectors[0]) if vectors else []
if not vec or (self._embed_dim and len(vec) != self._embed_dim):
raise ValueError(
f"{what} embedding failed: expected a non-empty "
f"{self._embed_dim or '?'}-dim vector, got {len(vec)} values."
)
return vec
def memorize(self, content: str, metadata: Mapping[str, object] | None = None) -> Memory:
"""Embed and persist a single memory; return the stored record.
Raises ``ValueError`` on empty content, an embedding-signature mismatch, or a failed
(empty/wrong-dim) embedding — never stores a corrupt vector.
"""
content = content.strip()
if not content:
raise ValueError("cannot memorize empty content")
meta = coerce_metadata(metadata) # validate JSON-safety before spending an embed call
vec = self._embed(content, "memory")
memory = Memory(
id=uuid.uuid4().hex,
user_id=self.user_id,
content=content,
created_at=_now_iso(),
metadata=meta,
)
# L5 anti-saturation — additive and OFF by default. Only when a flag is on AND the store
# supports supersessions do we scan the user's neighbours; otherwise this whole block is
# skipped and the path below is byte-identical to pre-L5. ``absorbed`` is a pre-existing
# memory the write folds into (dedup/NOOP → no new row); ``supersede_old_id`` is the id the
# new memory replaces (soft-update UPDATE → insert new row + record the supersession).
supersede_old_id: str | None = None
if (self.use_dedup or self.use_soft_update) and hasattr(self.store, "add_supersession"):
absorbed, supersede_old_id = self._antisaturation(memory, vec)
if absorbed is not None:
return absorbed
self.store.add(memory, vec)
if supersede_old_id is not None:
# UPDATE verdict: the new memory supersedes the old value so recall returns only the
# latest. Best-effort — a failure here (e.g. a stale/bad id) must never break memorize.
# ``add_supersession`` is an optional store capability (resolved via getattr so the
# engine types against the core ``MemoryStore`` Protocol); it is always present here
# because ``supersede_old_id`` is only set inside the ``hasattr`` branch above.
record = getattr(self.store, "add_supersession", None)
if callable(record):
try:
record(supersede_old_id, memory.id, memory.created_at)
except Exception:
log.warning(
"supersession record failed for memory %s; continuing", memory.id[:8]
)
# L4 episodic + G2 graph write-time enrichment — additive and OFF by default. BOTH features
# derive from ONE shared ``extractor.generate`` call (the folded flash-lite extraction), so
# a memorize with both on still spends a single aux call. When neither flag is on (or no
# extractor is configured) this whole block is skipped and the path above is byte-identical
# to the pre-enrichment engine. Any failure is swallowed: enrichment must NEVER make
# memorize fail or change the stored memory.
self._enrich(memory)
return memory
def _antisaturation(
self, memory: Memory, vec: Sequence[float]
) -> tuple[Memory | None, str | None]:
"""Decide how L5 should absorb ``memory`` before it is inserted.
Returns ``(absorbed, supersede_old_id)``:
- ``(existing, None)`` — a duplicate (dedup) or redundant restatement (arbiter NOOP): the
caller returns ``existing`` and inserts NOTHING, bounding store growth.
- ``(None, old_id)`` — a contradiction the arbiter ruled UPDATE: the caller inserts the new
row and records that ``old_id`` is now superseded (so recall returns only the new value).
- ``(None, None)`` — a plain ADD (unrelated, or no arbiter): the caller inserts normally.
"""
related = self._related_memories(vec)
if not related:
return None, None
# Fetch the superseded set ONCE. Both branches below must ignore superseded neighbours: a
# restated-but-currently-true fact whose nearest neighbour is an OLD (superseded) row must
# be written fresh, never absorbed back into the dead row (which would make it vanish from
# recall — the A→B→A "Paris→Berlin→Paris" loss).
superseded = self._superseded_ids()
if self.use_dedup:
# Absorb into the highest-cosine LIVE near-duplicate; if the only near-duplicate is a
# superseded row, fall through and write the new (currently-true) memory fresh.
dup = next(
(m for m, c in related if c >= self.dedup_threshold and m.id not in superseded),
None,
)
if dup is not None:
return dup, None # near-identical to a LIVE memory → skip the write
if self.use_soft_update and self.arbiter is not None:
candidate = next(
(
m
for m, c in related
if _SOFT_UPDATE_FLOOR <= c < self.dedup_threshold and m.id not in superseded
),
None,
)
if candidate is not None:
verdict, old_id = self._arbitrate(memory, candidate)
if verdict == SUPERSEDE_NOOP:
return candidate, None
if verdict == SUPERSEDE_UPDATE:
return None, old_id
return None, None
def _related_memories(self, vec: Sequence[float]) -> list[tuple[Memory, float]]:
"""Return this user's nearest existing memories to ``vec`` as ``(Memory, cosine)``, desc.
Reuses the store's existing ``build_index`` + dense search (works for both SQLite and
Postgres). At personal scale this brute-force scan is the same one recall already runs.
"""
index, by_id = self.store.build_index(self.user_id)
if len(index) == 0:
return []
out: list[tuple[Memory, float]] = []
for chunk in index.dense_search(vec, _NEIGHBOR_POOL):
mem = by_id.get(chunk.session_id)
if mem is not None:
out.append((mem, _cosine(vec, chunk.embedding)))
return out
def _arbitrate(self, memory: Memory, candidate: Memory) -> tuple[str, str | None]:
"""Run ONE cheap arbiter call to classify ``memory`` vs ``candidate``. NEVER raises.
Returns ``(verdict, supersede_old_id)``. On UPDATE the id is bound to the ``candidate`` we
actually showed the arbiter (never a free-form id it invents) — so a supersession can only
ever target a real, user-scoped memory. Any failure falls back to a plain ADD.
"""
arbiter = self.arbiter
if arbiter is None: # pragma: no cover — guarded by the caller; narrows for typing
return "ADD", None
try:
prompt = build_supersession_arbiter_prompt(
memory.content, candidate.id, candidate.content, memory.created_at
)
result = arbiter.generate(prompt, temperature=0.0, max_output_tokens=128)
parsed = parse_supersession_verdict(result.text)
verdict = parsed["verdict"] or "ADD"
# Bind the supersession to the candidate we actually showed the arbiter — never a
# free-form id it might invent — so a supersession can only target a real, user-scoped
# memory. (The parsed ``supersedes_id`` confirms intent; the write uses the candidate.)
return verdict, candidate.id
except Exception:
log.warning("soft-update arbiter failed for memory %s; adding plainly", memory.id[:8])
return "ADD", None
def _enrich(self, memory: Memory) -> None:
"""Run the ONE shared write-time extraction and fan out to episodic + graph (best-effort).
A single ``extractor.generate`` call feeds BOTH features. The block is skipped entirely when
no extractor is configured or neither feature is enabled/supported by the store — so
memorize is byte-identical to the pre-enrichment path. The generate call and each downstream
store are individually guarded: an extraction or persistence failure NEVER makes memorize
fail or change the stored memory. Graph fields are requested only when ``use_graph`` is on,
so the episodic-only path (``use_graph`` off) keeps its exact prompt + token budget.
"""
extractor = self.extractor
if extractor is None:
return
want_episodic = self.use_episodic and hasattr(self.store, "add_event")
want_graph = self.use_graph and hasattr(self.store, "ensure_self_entity")
if not (want_episodic or want_graph):
return
try:
prompt = build_episodic_extraction_prompt(
memory.content, memory.created_at, include_graph=want_graph
)
# 256 tokens (episodic-only) keeps the graph-off budget byte-identical; the graph
# payload (entities + relations) needs more room, so widen ONLY when graph is requested.
max_tokens = 512 if want_graph else 256
text = extractor.generate(prompt, temperature=0.0, max_output_tokens=max_tokens).text
except Exception:
log.warning("write-time extraction failed for memory %s; continuing", memory.id[:8])
return
if want_episodic:
self._store_event(memory, text)
if want_graph:
self._store_graph(memory, text)
def _store_event(self, memory: Memory, text: str) -> None:
"""Parse the episodic fields from the shared extraction and write the ``events`` row.
Best-effort: any parse/write error is logged and swallowed so episodic enrichment can never
fail memorize. Byte-identical to the pre-graph episodic write (same parser, same row).
"""
# ``add_event`` is an optional store capability (resolved via getattr so the engine types
# against the core ``MemoryStore`` Protocol); it is present here because ``_enrich`` only
# calls this method when ``hasattr(self.store, "add_event")``.
add_event = getattr(self.store, "add_event", None)
if not callable(add_event):
return
try:
parsed = parse_episodic_extraction(text, fallback_event_time=memory.created_at)
add_event(
memory_id=memory.id,
user_id=memory.user_id,
event_time=parsed["event_time"],
ingest_time=memory.created_at,
actor=parsed["actor"],
location=parsed["location"],
event_type=parsed["event_type"],
)
except Exception:
log.warning("episodic extraction failed for memory %s; continuing", memory.id[:8])
def _store_graph(self, memory: Memory, text: str) -> None:
"""Parse the graph half of the extraction and persist it for ``memory`` (best-effort).
Resolves the ego ``self`` root, upserts every extracted entity (self aliases collapse to the
self node and are NEVER a row), adds each labeled directed relation tagged with the source
``memory.id``, then links the memory to its subject entity (role=subject) and every other
mentioned entity (role=mention). Wrapped whole in try/except: a bad extraction must NEVER
fail memorize or touch the stored memory.
"""
# The entity-graph methods (ensure_self_entity/upsert_entity/add_entity_edge/
# link_memory_entity) are an OPTIONAL store capability outside the core ``MemoryStore``
# Protocol; ``_enrich`` only calls this method when ``hasattr(self.store,
# "ensure_self_entity")``. Bind the store as ``Any`` here so this duck-typed enrichment
# surface stays additive and never leaks into the engine's core typed contract.
store: Any = self.store
try:
parsed = parse_graph_extraction(text)
self_id = store.ensure_self_entity(self.user_id)
# norm_name -> entity_id cache (self aliases collapse to self_id, never a real row)
resolved: dict[str, str] = {}
# norm_name -> declared type, so a relation endpoint upserts with its stated type
declared: dict[str, str] = {}
for ent in parsed["entities"]:
norm = _norm_entity_name(ent["name"])
if norm and norm not in _SELF_ALIASES:
declared[norm] = ent["type"]
def resolve(raw: str) -> str | None:
norm = _norm_entity_name(raw)
if not norm:
return None
if norm in _SELF_ALIASES:
return self_id
eid = resolved.get(norm)
if eid is None:
etype = declared.get(norm, "thing")
eid = store.upsert_entity(self.user_id, raw.strip(), etype)
resolved[norm] = eid
return eid
mentioned: set[str] = set()
# upsert every declared entity (linked as a mention even without an explicit relation)
for ent in parsed["entities"]:
eid = resolve(ent["name"])
if eid is not None and eid != self_id:
mentioned.add(eid)
for rel in parsed["relations"]:
src_id, dst_id = resolve(rel["src"]), resolve(rel["dst"])
if src_id is None or dst_id is None:
continue
store.add_entity_edge(self.user_id, src_id, rel["label"], dst_id, memory.id)
mentioned.update(e for e in (src_id, dst_id) if e != self_id)
subject = parsed["subject"]
subject_id = self_id
if subject and _norm_entity_name(subject) not in _SELF_ALIASES:
subject_id = resolve(subject) or self_id
store.link_memory_entity(self.user_id, memory.id, subject_id, "subject")
for eid in mentioned:
if eid != subject_id:
store.link_memory_entity(self.user_id, memory.id, eid, "mention")
except Exception:
log.warning("graph extraction failed for memory %s; continuing", memory.id[:8])
def recall(self, query: str, limit: int | None = None) -> list[Memory]:
"""Return the memories most relevant to ``query`` (hybrid dense+lexical, RRF).
Embeds the query once; no generation. Returns ``[]`` for an empty query, an empty
store, or ``limit <= 0``.
"""
query = query.strip()
if not query:
return []
k = self.top_k if limit is None else limit
if k <= 0:
return []
qvec = self._embed(query, "query")
# With a reranker configured, over-retrieve a deeper POOL so the cross-encoder has real
# candidates to reorder, then filter superseded rows and rerank down to ``k``.
if self.reranker is not None:
pool = self._retrieve_pool(query, qvec, max(self.rerank_pool, k))
return self._rerank(query, self._filter_superseded(pool), k)
# No reranker AND no supersession filter: the byte-identical pre-L5 path — fetch exactly
# ``k`` and return them (the filter is a no-op that returns the pool unchanged, and it
# never consults the store's supersession set).
if not self._supersession_filter_active():
return self._filter_superseded(self._retrieve_pool(query, qvec, k))
# Supersession filter IS active: fetching only ``k`` would under-fill once the superseded
# rows are dropped, so over-fetch with headroom and grow the depth until we have ``k`` LIVE
# results (or the store is exhausted). Retrieval is RRF-ordered, so a deeper fetch only
# extends the tail — the top-``k`` live ordering is preserved.
depth = k + _SUPERSEDE_HEADROOM
while True:
pool = self._retrieve_pool(query, qvec, depth)
live = self._filter_superseded(pool)
if len(live) >= k or len(pool) < depth:
return live[:k]
depth *= 2
def _retrieve_pool(self, query: str, qvec: Sequence[float], depth: int) -> list[Memory]:
"""Return up to ``depth`` RRF-ordered candidate memories (superseded rows NOT yet filtered).
Prefers a store-side hybrid pushdown (Postgres/pgvector: dense ``<=>`` + FTS, fused in-DB) —
lower latency + no full-store load. Stores without it (SQLite) materialise the user's rows
and run the tested in-memory hybrid pipeline.
"""
search = getattr(self.store, "search", None)
if callable(search):
return list(search(self.user_id, query, qvec, depth))
index, by_id = self.store.build_index(self.user_id)
if len(index) == 0:
return []
out: list[Memory] = []
for chunk in hybrid_retrieve(index, query, qvec, depth):
mem = by_id.get(chunk.session_id)
if mem is not None:
out.append(mem)
return out
def _rerank(self, query: str, pool: list[Memory], k: int) -> list[Memory]:
"""Reorder an RRF-ordered candidate ``pool`` with the cross-encoder, keeping the top ``k``.
With no reranker configured this returns ``pool`` UNCHANGED — the load-bearing guard for the
byte-identical no-reranker path (recall neither over-fetches nor reorders). When a reranker
is set, it scores every ``(id, content)`` candidate and returns the top ``k`` in its order;
any ids it drops are backfilled in RRF order so a short result still yields a full ``k``. On
``VertexRerankError`` or ANY exception we log and fall back to the pool's RRF top-``k``
(``pool[:k]``) — recall must never fail because of reranking.
"""
if self.reranker is None or not pool:
return pool
try:
ranked_ids = self.reranker.rerank(query, [(m.id, m.content) for m in pool], top_k=k)
except Exception: # noqa: BLE001 — VertexRerankError or anything else → graceful fallback
log.warning("rerank failed for user %s; falling back to RRF top-%d", self.user_id, k)
return pool[:k]
by_id = {m.id: m for m in pool}
# Dedupe the returned ids preserving order (seen-set): a misbehaving reranker that repeats
# an id must NOT inject the same Memory twice into the result.
seen: set[str] = set()
ranked: list[Memory] = []
for i in ranked_ids:
if i in by_id and i not in seen:
seen.add(i)
ranked.append(by_id[i])
if len(ranked) < k: # backfill dropped/unknown ids in the original RRF order (no dupes)
for m in pool:
if m.id not in seen:
seen.add(m.id)
ranked.append(m)
return ranked[:k]
def _filter_superseded(self, memories: list[Memory]) -> list[Memory]:
"""Drop superseded memories so recall returns the LATEST value — GATED on anti-saturation.
The FIRST guard is load-bearing for the byte-identical guarantee: with BOTH L5 flags off we
return the input list UNCHANGED and NEVER query ``superseded_ids`` — recall is exactly the
pre-L5 result. Only when dedup/soft-update is active (and the store supports it) do we
consult the supersession set and filter.
"""
if not self._supersession_filter_active():
return memories
superseded = self._superseded_ids()
if not superseded:
return memories
return [m for m in memories if m.id not in superseded]
def _supersession_filter_active(self) -> bool:
"""True iff anti-saturation is on AND the store exposes the supersession set.
Load-bearing for the byte-identical guarantee: with BOTH L5 flags off this is False, so
recall neither over-fetches nor consults the supersession set — exactly the pre-L5 path.
"""
return (self.use_dedup or self.use_soft_update) and hasattr(self.store, "superseded_ids")
def _superseded_ids(self) -> set[str]:
"""This user's superseded memory ids, or an empty set if the store predates the feature.
``superseded_ids`` is an optional store capability, resolved via getattr so the engine can
type against the core ``MemoryStore`` Protocol without the enrichment surface leaking in.
"""
fn = getattr(self.store, "superseded_ids", None)
if not callable(fn):
return set()
result: set[str] = fn(self.user_id)
return result
def list_memories(self, limit: int = 20) -> list[Memory]:
"""Return the most recently added memories first."""
return self.store.list_recent(self.user_id, limit)
def timeline(
self, since: str | None = None, until: str | None = None, limit: int = 50
) -> list[Memory]:
"""Return this user's episodic memories ordered by event_time (L4; additive).
Delegates to ``store.timeline``. Returns ``[]`` for a store that predates episodic
support, so callers never break on the pre-L4 store surface.
"""
timeline_fn = getattr(self.store, "timeline", None)
if not callable(timeline_fn):
return []
results: list[Memory] = timeline_fn(self.user_id, since=since, until=until, limit=limit)
return results
def forget(self, memory_id: str) -> bool:
"""Delete one memory by exact id; return True if it existed."""
return self.store.delete(memory_id, self.user_id)
def forget_prefix(self, memory_id: str) -> str:
"""Delete by full id OR an unambiguous short-id prefix; return a status message.
Refuses to delete when a short prefix matches MULTIPLE memories — prevents a short id
from silently deleting the wrong memory (data loss). Read-only resolution is SQL-side
and bounded (no full-table scan).
"""
target = memory_id.strip()
if not target:
return "No memory id given."
if len(target) < 32: # a short id (full uuid4 hex is 32 chars)
matches = self.store.resolve_id_prefix(self.user_id, target, limit=2)
if not matches:
return f"No memory matched {memory_id!r}."
if len(matches) > 1:
return (
f"Ambiguous id {memory_id!r}: it matches multiple memories — "
"use the full id from list_memories."
)
target = matches[0]
ok = self.forget(target)
return f"Forgot {target[:8]}." if ok else f"No memory matched {memory_id!r}."
def count(self) -> int:
"""Total memories stored for this user."""
return self.store.count(self.user_id)
def close(self) -> None:
"""Close the underlying store (flush SQLite); safe to call once at shutdown."""
self.store.close()