Skip to main content

Graph-based memory system for AI agents — cognitive science meets deep learning

Project description

memorygraph

Graph-based memory system for AI agents — cognitive science meets deep learning.

memorygraph provides a production-ready, research-grade memory layer for LLM agents, autonomous systems, and any AI application that needs to store, retrieve, and reason over structured knowledge across time.

Python 3.9+ License: MIT PyTorch optional


Overview

Most agent memory systems are flat key-value stores or simple vector databases. memorygraph treats memory the way cognitive science does — as a dynamic, associative graph where:

  • Nodes are memories (semantic facts, episodic events, procedural skills, emotional anchors)
  • Edges are typed relationships (causal, hierarchical, associative, temporal, similarity)
  • Activation spreads through the graph like priming in the human brain
  • Forgetting follows empirically validated curves (Ebbinghaus, Power Law, FSRS-4.5)
  • Working memory enforces a bounded attentional buffer with interference effects
  • Neural components (Modern Hopfield, DNC, MemN2N, Transformer encoder) are available for learned retrieval

Architecture

memorygraph/
├── core/
│   ├── types.py          # Enums, protocols, exception hierarchy, type aliases
│   ├── node.py           # MemoryNode — ACT-R activation, retention models, embeddings
│   ├── edge.py           # MemoryEdge — Hebbian weight updates, typed relationships
│   └── graph.py          # MemoryGraph — thread-safe CRUD, search, graph algorithms
│
├── memory/
│   └── working.py        # Baddeley working memory buffer (phonological, visuospatial,
│                         #   episodic, central executive slots)
│
├── retrieval/
│   └── spreading.py      # Collins & Loftus spreading activation (Dijkstra-based)
│
├── decay/
│   └── forgetting.py     # Forgetting models: Ebbinghaus, Power Law, SM-2, FSRS-4.5
│                         # SpacedRepetitionScheduler with per-node review tracking
│
├── neural/               # PyTorch components (optional — graceful degradation)
│   ├── hopfield.py       # Modern Hopfield Networks (Ramsauer et al., 2020)
│   ├── encoder.py        # Transformer memory encoder + contrastive loss
│   └── retriever.py      # DNC-style differentiable memory + End-to-End Memory Networks
│
└── persistence/
    └── storage.py        # JSON, Pickle, SQLite backends; backup rotation; npz export

Installation

Minimal (NumPy only):

pip install memorygraph

With PyTorch neural components:

pip install "memorygraph[torch]"

Full installation:

pip install "memorygraph[full]"

Development:

git clone https://github.com/memorygraph/memorygraph
cd memorygraph
pip install -e ".[dev,torch]"

Quick Start

Basic Memory Graph

import numpy as np
from memorygraph import (
    MemoryGraph, MemoryNode, MemoryEdge,
    MemoryType, EdgeType,
)

# Create a graph with 768-dimensional embeddings (e.g. BERT output size)
graph = MemoryGraph(embedding_dim=768, name="agent_memory")

# Add a semantic memory node
fact = MemoryNode(
    content      = "The Eiffel Tower is located in Paris, France.",
    memory_type  = MemoryType.SEMANTIC,
    embedding    = np.random.randn(768).astype(np.float32),  # use your encoder
    importance   = 0.85,
    tags         = ["geography", "landmarks", "europe"],
)
graph.add_node(fact)

# Add an episodic memory
event = MemoryNode(
    content           = "User asked about European landmarks at 14:32.",
    memory_type       = MemoryType.EPISODIC,
    embedding         = np.random.randn(768).astype(np.float32),
    importance        = 0.6,
    emotional_valence = 0.2,   # slightly positive
)
graph.add_node(event)

# Link them causally
graph.add_edge(MemoryEdge(
    source_id = event.id,
    target_id = fact.id,
    edge_type = EdgeType.CAUSAL,
    weight    = 0.9,
))

print(graph.stats())

Embedding-Based Retrieval

query_embedding = np.random.randn(768).astype(np.float32)

results = graph.similarity_search(
    query     = query_embedding,
    k         = 10,
    threshold = 0.5,
    memory_type = MemoryType.SEMANTIC,  # optional filter
    tags      = ["geography"],          # optional tag filter (AND)
)

for result in results:
    node = graph.get_node(result.node_id)
    print(f"[{result.score:.4f}] {node.content}")

Spreading Activation

Retrieve memories by simulating associative priming — activation flows through the graph from seed nodes, decaying with distance and modulated by edge weights and node importance.

from memorygraph import SpreadingActivation, SpreadingConfig

config = SpreadingConfig(
    max_depth          = 4,
    activation_decay   = 0.6,
    use_edge_weights   = True,
    use_node_importance = True,
    use_temporal_boost = True,   # recently accessed nodes get a boost
    return_top_k       = 20,
    sort_by            = "activation",
)

spreader = SpreadingActivation(graph, config)
activated = spreader.spread(source_ids=[fact.id])

for result in activated:
    node = graph.get_node(result.node_id)
    print(f"[act={result.activation:.4f} depth={result.depth}] {node.content}")

# Cognitive priming: does seeing A make B more accessible?
primed = spreader.priming_search(prime_ids=[event.id], query_id=fact.id)

# Find the conceptual chain between two distant memories
chain = spreader.concept_chain(start_id=node_a.id, end_id=node_b.id)

Working Memory Buffer

A bounded attentional buffer modelled after Baddeley's multi-component working memory model. Supports four slot types with independent capacities, interference between similar items, and LRU/activation-based eviction.

from memorygraph import WorkingMemoryBuffer, WorkingMemorySlot

wm = WorkingMemoryBuffer(
    capacity          = 7,     # Miller's magical number
    decay_rate        = 0.05,
    interference_rate = 0.1,
    eviction_policy   = "lru_activation",
)

# Push nodes into specific slots
wm.push(node,  slot=WorkingMemorySlot.CENTRAL,       activation=1.0, priority=1.0)
wm.push(fact,  slot=WorkingMemorySlot.PHONOLOGICAL,  activation=0.8)
wm.push(event, slot=WorkingMemorySlot.EPISODIC,      activation=0.7)

# Tick to apply decay (call periodically or on each agent step)
wm.tick()

# Inspect the buffer
snapshot = wm.snapshot()
print(f"Utilization: {snapshot['utilization']*100:.0f}%")

# Get the current focus of attention
focus_item = wm.focus()

# Register eviction callback
wm.on_evict = lambda item: print(f"Evicted: {item.node.content}")

Spaced Repetition & Forgetting

Schedule memory reviews using FSRS-4.5, SM-2, Ebbinghaus, or Power Law models. Track stability and difficulty per node.

from memorygraph import SpacedRepetitionScheduler

scheduler = SpacedRepetitionScheduler(model="fsrs", target_retention=0.90)

# Register nodes for spaced repetition
for node in graph.query_by_type(MemoryType.SEMANTIC):
    scheduler.register_node(node.id, initial_grade=node.importance)

# Record a review (grade: 0.0 = forgot, 1.0 = perfect recall)
entry = scheduler.record_review(node_id=fact.id, grade=0.9)
print(f"Next review in {entry.time_until_due() / 86400:.1f} days")
print(f"Stability: {entry.stability:.2f}")

# Get all nodes due for review right now
due_nodes = scheduler.get_due_nodes()

# Get upcoming reviews within the next 3 days
upcoming = scheduler.get_upcoming(within_seconds=3 * 86400)

# Plot the forgetting curve for a specific node
t_array, r_array = scheduler.forgetting_curve_points(
    node_id          = fact.id,
    num_points       = 100,
    time_horizon_days = 30,
)

Persistence

from memorygraph import GraphStorageManager

manager = GraphStorageManager(backend="sqlite")  # or "json", "pickle"
manager.save(graph, "agent_memory.db")

# Load back
graph = manager.load("agent_memory.db")

# Automatic backup rotation (keeps last 3 backups)
manager.save_with_backup("agent_memory.db", keep_backups=3)

# Periodic checkpoints
manager.checkpoint(directory="./checkpoints", prefix="agent")

# Export / import embeddings separately (for fine-tuning pipelines)
manager.export_embeddings(graph, "embeddings.npz")
manager.import_embeddings(graph, "embeddings_v2.npz")

Transactions

with graph.transaction():
    graph.add_node(node_a)
    graph.add_node(node_b)
    graph.add_edge(edge)
    # If any operation raises an exception, the entire transaction rolls back

Neural Components (PyTorch)

All neural components require pip install "memorygraph[torch]". They are imported lazily — if PyTorch is not installed, the rest of the library works without modification.

Modern Hopfield Networks

Based on Ramsauer et al. (2020). Exponential storage capacity compared to classical Hopfield networks. Mathematically equivalent to attention with a specific energy function.

from memorygraph.neural import ModernHopfieldLayer, HopfieldMemoryPool, HopfieldStack
import torch

layer = ModernHopfieldLayer(
    input_dim  = 512,
    hidden_dim = 512,
    num_heads  = 8,
    beta       = 4.0,   # inverse temperature (higher = sharper retrieval)
    num_iter   = 3,     # retrieval iterations
)

query    = torch.randn(batch, seq_q, 512)
patterns = torch.randn(batch, seq_k, 512)
output, attn_weights = layer(query, patterns, return_attn=True)

# Energy of the Hopfield network
energy = layer.energy(query, patterns)

# Fixed learnable memory pool
pool = HopfieldMemoryPool(capacity=1024, pattern_dim=512, num_heads=8)
retrieved, weights = pool(query)

# Deep stacked Hopfield with geometric beta schedule
stack = HopfieldStack(
    input_dim     = 512,
    num_layers    = 6,
    beta_schedule = "geometric",
    initial_beta  = 1.0,
    final_beta    = 8.0,
)
output, all_attentions = stack(query, patterns, return_all_attn=True)

Transformer Memory Encoder

from memorygraph.neural import MemoryEncoder, MemoryContrastiveLoss
import torch

encoder = MemoryEncoder(
    vocab_size  = 30522,   # BERT vocabulary
    dim         = 768,
    num_heads   = 12,
    num_layers  = 6,
    max_len     = 512,
    pooling     = "mean",  # or "cls", "max", "attn"
)

input_ids = torch.randint(0, 30522, (batch, seq_len))
output = encoder(input_ids)

print(output.embedding.shape)   # (batch, 768)
print(output.importance)        # (batch, 1)  — learned importance score
print(output.emotion)           # (batch, 1)  — emotional valence
print(output.confidence)        # (batch, 1)  — retrieval confidence

# Contrastive training with InfoNCE / NT-Xent loss
loss_fn = MemoryContrastiveLoss(temperature=0.07)
labels  = torch.arange(batch)
loss    = loss_fn(output.embedding, output.embedding, labels)
loss.backward()

Differentiable Neural Memory (DNC-inspired)

from memorygraph.neural import NeuralMemoryModule
import torch

dnc = NeuralMemoryModule(
    memory_size    = 128,   # number of memory slots
    memory_dim     = 64,    # dimensionality per slot
    controller_dim = 512,   # controller hidden size
    num_reads      = 4,     # parallel read heads
)

# Process a sequence; state carries memory across steps
controller_sequence = torch.randn(batch, time_steps, 512)
output_sequence, final_state = dnc.process_sequence(controller_sequence)
# output_sequence: (batch, time_steps, 512)  controller enriched with memory reads
# final_state.memory: (batch, 128, 64)       memory after all writes

End-to-End Memory Networks (MemN2N)

Based on Sukhbaatar et al. (2015). Multi-hop reasoning over a bag of memory sentences.

from memorygraph.neural import MemoryNetwork
import torch

model = MemoryNetwork(
    vocab_size   = 10000,
    embed_dim    = 256,
    num_hops     = 3,
    output_size  = 10000,   # answer vocabulary
)

query_tokens  = torch.randint(0, 10000, (batch, q_len))
memory_tokens = torch.randint(0, 10000, (batch, num_memories, sent_len))

logits = model(query_tokens, memory_tokens)   # (batch, output_size)

Core Concepts

Memory Node Lifecycle

Register → Encode → Store → Retrieve → Reinforce / Decay → Consolidate → (Forget)

Each MemoryNode tracks:

Attribute Description
importance 0–1 base importance score
emotional_valence –1 to +1 emotional tone
confidence 0–1 retrieval confidence
consolidation_state LABILE → CONSOLIDATING → CONSOLIDATED → LONG_TERM
temporal_info access history, ACT-R activation computation
embedding dense vector representation
tags string tags for structured filtering

ACT-R Base-Level Learning:

$$A_i = \ln\left(\sum_j t_j^{-d}\right) + \alpha \cdot e_i \cdot c_i$$

where $t_j$ are times since past accesses, $d$ is the decay parameter, $\alpha$ is emotional boost, $e_i$ is emotional valence, and $c_i$ is confidence.

Retention Models:

Model Formula Use Case
Ebbinghaus $R = e^{-t/S}$ Classic forgetting curve
Power Law $R = (1 + bt)^{-\alpha}$ Power forgetting (Wixted & Ebbesen)
Hyperbolic $R = 1/(1 + t/S)$ Hyperbolic discounting
FSRS-4.5 Spaced repetition algorithm Anki-compatible scheduling

Edge Types

Type Meaning
ASSOCIATION General semantic association
CAUSAL A caused B
TEMPORAL A preceded B
HIERARCHICAL A is a kind of / part of B
SIMILARITY A and B are similar
CONTRADICTS A contradicts B
COMPOSITIONAL A is composed of B
EMOTIONAL A triggered the emotion in B
EPISODIC_LINK Same episode / context
INFERENCE B was inferred from A

Consolidation States

LABILE  →  CONSOLIDATING  →  CONSOLIDATED  →  LONG_TERM
  ↑                                                ↓
  └──────────────── RECONSOLIDATING ───────────────┘

Call graph.advance_consolidation(node_id) to progress a node through the consolidation pipeline. Long-term memories are protected from prune_forgotten().


Graph Algorithms

# Breadth-first search up to depth 3, only CAUSAL edges
nodes = graph.bfs(start_id, max_depth=3, edge_type=EdgeType.CAUSAL)

# Dijkstra shortest path (cost = 1 - edge_weight)
path = graph.shortest_path(source_id, target_id)

# Strongly connected components (Kosaraju's algorithm)
sccs = graph.strongly_connected_components()

# Merge near-duplicate nodes (cosine similarity > threshold)
merged = graph.merge_similar_nodes(similarity_threshold=0.92)

# Remove forgotten memories below retention threshold
pruned = graph.prune_forgotten(retention_threshold=0.1)

Thread Safety

MemoryGraph and WorkingMemoryBuffer are fully thread-safe. All mutating operations acquire a reentrant lock (threading.RLock). The transaction() context manager holds the lock for the duration of the block.


Benchmarks

Tested on a consumer laptop (Intel Core i7-12700H, 32 GB RAM) with NumPy embeddings (no GPU):

Operation N nodes Time
add_node 10 000 < 1 ms / node
similarity_search (cosine, k=20) 10 000 ~12 ms
similarity_search (cosine, k=20) 100 000 ~110 ms
spreading_activation (depth=4) 10 000 ~8 ms
bfs (depth=3) 10 000 ~2 ms
SQLite save 10 000 ~1.4 s
SQLite load 10 000 ~0.9 s

For sub-millisecond ANN search at scale, install faiss-cpu and swap the embedding index.


Design Decisions

Why a graph and not a vector database?
Vector databases excel at "find the 10 most similar things." A graph additionally encodes why memories are related, supports reasoning along paths, models causal and temporal structure, and enables spreading activation — the associative priming that makes retrieval feel natural rather than purely similarity-driven.

Why is PyTorch optional?
The graph core (node CRUD, search, spreading activation, working memory, forgetting models, persistence) has zero heavy dependencies. An agent can run the full cognitive architecture on a Raspberry Pi. PyTorch components are addons for learned retrieval and training pipelines.

Why FSRS-4.5 and not just Ebbinghaus?
FSRS-4.5 is the current state of the art in spaced repetition, used by millions of Anki users. Its stability and difficulty parameters are empirically validated. Ebbinghaus and Power Law are provided for research comparisons and lightweight deployments.

Thread safety via RLock vs. actor model?
An agent's memory is typically accessed by one or two threads (main loop + background consolidation). RLock is simpler, has negligible overhead at this concurrency level, and avoids the latency of message-passing queues.


References

  • Ramsauer et al. (2020). Hopfield Networks is All You Need. ICLR 2021. arXiv:2008.02217
  • Sukhbaatar et al. (2015). End-To-End Memory Networks. NeurIPS 2015. arXiv:1503.08895
  • Graves et al. (2016). Hybrid computing using a neural network with dynamic external memory (DNC). Nature 538.
  • Anderson et al. (2004). An integrated theory of the mind (ACT-R). Psychological Review.
  • Baddeley (2000). The episodic buffer: a new component of working memory. Trends in Cognitive Sciences.
  • Collins & Loftus (1975). A spreading-activation theory of semantic processing. Psychological Review.
  • Ebbinghaus (1885). Über das Gedächtnis.
  • Wixted & Ebbesen (1991). On the form of forgetting. Psychological Science.
  • Ye et al. (2022). A New Algorithm for Optimizing Spaced Repetition Scheduling (FSRS). arXiv:2402.10340

License

MIT — see LICENSE.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

neural_memorygraph-1.0.0.tar.gz (65.6 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

neural_memorygraph-1.0.0-py3-none-any.whl (62.4 kB view details)

Uploaded Python 3

File details

Details for the file neural_memorygraph-1.0.0.tar.gz.

File metadata

  • Download URL: neural_memorygraph-1.0.0.tar.gz
  • Upload date:
  • Size: 65.6 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.14.4

File hashes

Hashes for neural_memorygraph-1.0.0.tar.gz
Algorithm Hash digest
SHA256 c414ad87a33866de5d9755ec81a95f724c28e89d90fbef14c8695776c5d9566f
MD5 2ef2b28dd84ddc10dd53355c1a4bac96
BLAKE2b-256 3a09dd37ec2c42fe98be4a9ccabf8118d92aa5a32c2ffbfd1c466df8a8f52bd6

See more details on using hashes here.

File details

Details for the file neural_memorygraph-1.0.0-py3-none-any.whl.

File metadata

File hashes

Hashes for neural_memorygraph-1.0.0-py3-none-any.whl
Algorithm Hash digest
SHA256 212d3ad688fc2c3fff109513d35c56e666458abdc9b53401e0217f61c7849454
MD5 cc0fa2e74ecf3ae0b039f6f0eab67585
BLAKE2b-256 b86a66dd94aaad8ee71b9e5eb5ea487afe1d16f1a3ead84b432855b99bbbc04c

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page