kgmm — KG-Markov Anticipatory Memory

A memory system where knowledge-graph nodes double as Markov states. Transition edges learned from session history predict the next state, so context is prefetched before it is asked for. Existing memory systems (Zep, Mem0, Supermemory) are reactive — retrieval starts when the query arrives. This adds a push layer.

Target use case: multi-agent pipelines. Every agent handoff normally costs a blocking retrieval round-trip. With anticipation, retrieval for the next agent happens while the upstream agent is still running — a throughput multiplier, plus token control (only top-k probable states get context, within a budget).

Quick start

python3 examples/multi_agent_demo.py     # no dependencies, stdlib only

Sample output:

agent      state                  retrieval      latency  tokens
------------------------------------------------------------------
Planner    django_models          cold fetch       250ms       9
           └─ prefetching for: migration_conflict (p=1.00)
Coder      migration_conflict     prefetch HIT       2ms       9
...
blocking retrieval latency :    1000ms (reactive baseline)
with anticipatory prefetch :     256ms (74% reduction)
context tokens per handoff : 192 (naive full-graph dump) vs 8 avg (budgeted)

API

from kgmm import Memory

memory = Memory()                                  # coding-domain ontology
result = memory.update("hit a migration conflict") # retrieve + learn + prefetch
result.primary                  # current Markov state: "migration_conflict"
result.context                  # context bundle for the current state
result.context_was_prefetched   # True if the previous prediction was right
result.predictions              # top-k likely next states with probabilities
result.prefetched               # bundles already built for those states
memory.end_session()            # write the session's transitions to the graph
memory.store.save("graph.json") # persist; GraphStore.load() to restore

Project layout

kgmm/        core library, shared by all versions
             (llm_extraction.py is the v3 extraction layer)
tests/       unit tests for the core
examples/    v1 multi-agent prefetch demo (stdlib-only)
eval/        experiments, organized by version — see eval/README.md
  v1/        rule-based system: prediction, latency, quality evals
  v2/        Graphiti/Neo4j content backend
  v3/        LLM extraction + live end-to-end pipeline
paper/       full research writeup (LaTeX, all results)

Architecture

message → extraction → entity resolution → primary node (current state)
                                                │
                              ┌─────────────────┴──────────────┐
                              ▼                                ▼
                     consume prefetch cache            Markov forward walk
                     (hit → ~0ms retrieval)         top-k next states, p ≥ 0.4
                                                               │
                                                               ▼
                                                   prefetch context bundles
                                                   (token-budgeted, by p)

module	role
ontology.py	fixed ~55-node coding state space (v1: closed ontology → consistent extraction)
extraction.py	RuleBasedExtractor (default, deterministic) / LLMExtractor (optional, claude-haiku)
resolution.py	entity → canonical node; threshold 0.85 is the key hyperparameter
graph.py	one node set, two edge types (semantic + transition); decay; JSON persistence
transitions.py	session → transition counts; blended personal/global cold-start prior
simulation.py	forward walk, top-k, optional depth-2 path probabilities
prefetch.py	budgeted context bundles, cache, hit/miss/waste stats
memory.py	the Memory facade tying it together (30-min session boundary)

Design decisions (per the brief)

• Domain-specific v1: fixed coding ontology, not open extraction.
• Cold start: p = λ·p_personal + (1−λ)·p_global, λ = min(1, sessions/20).
• Confidence gate: prefetch only when p ≥ 0.4 — wrong predictions must not poison context.
• Decay: effective_count = count · exp(−λ_decay · days) — workflows change.
• Top-k = 3 predicted states, token budget allocated proportional to probability.

v2 path (not yet wired)

• LLMExtractor — open extraction; needs pip install kgmm[llm] + ANTHROPIC_API_KEY.
• Neo4j/Graphiti backend — GraphStore's interface (nodes, typed edges, count/probability/last_seen as edge properties) maps 1:1 onto Neo4j; Graphiti replaces extraction.py + resolution.py + semantic retrieval, and the transition layer stays as-is on top.
• Embedding-based resolution — pass an embedder to EntityRegistry.

Tests

uv run --with pytest python -m pytest