Mnemosyne

A memory persistent ai agent

Comment

Mnemosyne

A Graph Neural Network Memory System for Persistent AI Agents

Inspiration

From the competition brief, we were given a list of problems to choose from. I read through all of them. Persistent memory was the one I kept coming back to.

Not because it seemed easy but because it is one of the harder problems on the list. But because it felt like the most honest problem. Every other capability you might build for an AI agent eventually runs into the same wall: the agent does not remember you. It does not know what you told it last week, what you prefer, what you are working on, or what you already tried. You rebuild context from scratch every session, and the agent stays a stranger no matter how many times you talk to it.

I chose persistent memory because I wanted to build something that solved that at a level deeper than storing a transcript and replaying it. I wanted to build a system where memory has structure, weight, and intelligence, where the agent does not just recall what you said, but understands what matters, what is related, and what can safely be forgotten.

The architectural idea that made this feel possible came from an unexpected direction. I had spent a lot of time observing how social media platforms, Twitter in particular, build personalized timelines for users from the very first interaction. What struck me was that the recommendation engine is not just matching content to stated preferences. It is learning a graph of relationships: what topics connect to what other topics, which accounts cluster together, how engagement on one piece of content predicts engagement on another. The platform is reasoning over a graph of the user's behavior, not just a flat list of their likes.

I sat with that observation for a while and asked a direct question: if an AI agent's memories are nodes, and semantic similarity creates edges between them, could a Graph Neural Network learn which memories actually matter for a given query better than a vector search alone can? That hypothesis is what Mnemosyne was built to test.

The name came naturally. I am an avid reader, and the name Mnemosyne had stayed with me since I first encountered it in the Odyssey. Mnemosyne is the goddess of memory in Greek mythology, the mother of the nine Muses, the keeper of all that is known and remembered. When it came time to name the project, the name came up on its own. It fit too well to use anything else.

What I Learned

Persistent memory in production is a largely unsolved problem. Before writing a single line of code I spent time researching how memory is actually handled in AI applications today. What I found was that most production systems either append the full conversation history to every prompt, which hits context window limits quickly and buries important information in noise, or store a rolling summary that loses specificity over time. Neither approach gives the agent genuine long-term memory. The gap between what the field describes as "memory" and what a genuinely persistent agent would require was wider than I expected.

The taxonomy of AI memory types. My research surfaced a distinction that became central to Mnemosyne's design: not all memories are the same kind of thing. There is a meaningful difference between a behavioral rule ("always give concise answers"), a stable fact ("the user is a software engineer"), a stated preference ("the user prefers dark mode"), and a past episode ("the user asked about Python async IO last Tuesday"). These types have different lifespans, different importance weights, and different roles in how an agent should reason. Building this taxonomy into the schema from the start, rather than treating all memories as undifferentiated text, was one of the most consequential early decisions in the project.

LangChain and LangGraph exist, and I chose not to use them. During my research I came across LangChain and LangGraph the two are widely used frameworks for building LLM applications, both of which include memory abstractions. I studied them carefully. LangChain offers conversation buffer memory, summary memory, and vector-store-backed memory. LangGraph adds stateful agent loops. What I found was that these are convenience abstractions built on top of simple patterns: append history, summarize it, or retrieve by similarity. None of them offer typed memory, semantic graph structure, learned relevance scoring, or principled forgetting. They are useful tools for shipping fast. They are not memory infrastructure. I chose to build from first principles because the problem required it, and because I wanted to understand every component rather than inheriting behavior I could not see into.

GNNs had not been applied to agent memory this way before. Before committing to the GraphSAGE approach I searched extensively for prior work in academic papers and open-source projects, on using graph neural networks specifically for persistent AI agent memory retrieval. I found work on knowledge graphs for question answering, graph-based recommendation systems, and GNN applications in NLP generally, but nothing that applied learned graph-based scoring to the problem of personal memory retrieval for a conversational agent. That absence was both a risk and an opportunity. It meant there was no established playbook to follow, but it also meant that if the hypothesis worked, it would be genuinely novel.

The social media analogy is deeper than it first appears. The more I worked on the retrieval problem, the more the Twitter parallel held up. Both systems face the same core challenge: given a large and growing set of items connected by a relational structure, rank the ones that are most relevant to this user right now. Twitter solves it with a graph of engagement signals. Mnemosyne solves it with a graph of semantic similarity and access patterns, scored by a GNN trained on relevance labels. The cold start problem is even structurally identical, a new user has no history, so the system cannot yet personalize, and must degrade gracefully to a safe default while real signal accumulates.

Graceful degradation is a design discipline. Every learned component in the system, the GNN relevance scores, the cluster confidence, must degrade gracefully to a safe baseline when data is sparse. A system that fails hard at cold start is not a memory system; it is a liability. Designing the confidence gate so the GNN always runs but its output is only trusted proportionally to the evidence behind it was one of the most important architectural decisions in the project.

Memory is a UX problem as much as a technical one. A system that generates a fresh random user_id every session structurally prevents memory accumulation regardless of how well the rest of the pipeline works. The most impactful single change in the entire project was replacing dynamic user id generation with a fixed persistent identifier. Architecture is nothing without coherent identity across sessions. This helped me in testing to ensure that the entire system was working as intended.

How I Built It

Mnemosyne is built around a four-layer pipeline that processes every conversational turn end-to-end.

Memory Schema

Every piece of information the agent learns is classified into one of four typed memory categories, each with a distinct role in the system:

Type Role Base Importance
RULE Behavioral directives the agent must follow $0.85$
PREFERENCE How the user likes things done $0.70$
FACT Stable information about the user $0.60$
EPISODE Specific past events $0.40$

Importance scores are adjusted at write time using a keyword boost:

$$\text{importance} = \min\left(1.0,\ \text{base_score} + \sum_{k \in \text{keywords}} 0.03 \cdot \mathbf{1}[k \in \text{content}]\right)$$

The Memory Graph

After every write, the new memory is compared against the user's existing memories in Qdrant. Any pair with cosine similarity at or above a fixed threshold receives a weighted edge in the graph:

$$\text{edge}(u, v) \iff \cos(\mathbf{e}_u, \mathbf{e}_v) \geq 0.75$$

This produces semantic clusters, memories about work connect to other work memories, preferences cluster with preferences, giving the GNN a structured graph to reason over rather than an unconnected set of points.

The GraphSAGE Model

The GNN takes each memory node as a $391$-dimensional feature vector:

$$\mathbf{x}i = [\underbrace{\mathbf{e}_i}{384} \mid \underbrace{\mathbf{t}i}{4} \mid \underbrace{s_i}{1} \mid \underbrace{a_i}{1} \mid \underbrace{c_i}_{1}]$$

where $\mathbf{e}_i$ is the sentence embedding, $\mathbf{t}_i$ is the one-hot memory type, $s_i$ is the importance score, $a_i$ is normalized age, and $c_i$ is normalized access count.

Two GraphSAGE layers with mean aggregation produce enriched node embeddings:

$$\mathbf{h}i^{(l+1)} = \sigma\left(\mathbf{W}^{(l)} \cdot \text{CONCAT}\left(\mathbf{h}_i^{(l)},\ \text{MEAN}{j \in \mathcal{N}(i)} \mathbf{h}_j^{(l)}\right)\right)$$

The model has two parallel output heads:

  • Relevance head: $r_i = \sigma(\text{MLP}(\mathbf{h}_i)) \in [0, 1]$ — predicts whether this memory will be accessed
  • Cluster head: $\mathbf{c}_i = \text{softmax}(\text{MLP}(\mathbf{h}_i)) \in \mathbb{R}^4$ — predicts memory type from neighborhood structure

Training uses a combined loss with dual supervision:

$$\mathcal{L} = \mathcal{L}{\text{BCE}}(r_i, y_i^{\text{rel}}) + \alpha \cdot \mathcal{L}{\text{CE}}(\mathbf{c}_i, y_i^{\text{type}}), \quad \alpha = 0.5$$

The cluster head acts as a regularizer, preventing the shared embedding from collapsing to shortcuts that ignore memory type structure.

Hybrid Retrieval

At inference time, four signals are combined into a single score per candidate memory:

$$\text{score}_i = 0.35 \cdot \text{sim}_i + 0.35 \cdot r_i + 0.15 \cdot \text{recency}_i + 0.15 \cdot \max(\mathbf{c}_i)$$

where recency decays smoothly with turns since last access:

$$\text{recency}_i = \frac{1}{1 + \delta_i \cdot 0.1}$$

Results are cached in Redis with a 30-second TTL, short enough that memory state changes are reflected quickly, long enough to absorb repeated queries within a turn.

Memory Aging and Forgetting

Every turn increments a turns_since_access counter for every active memory belonging to the user. Accessed memories have this counter reset to zero. Every ten turns, a forgetting pass archives memories whose importance and recency have both decayed below threshold. Archived memories are never deleted, only excluded from active retrieval, preserving the option to recover or audit historical memory without data loss.

The LLM Integration

Retrieved memories are grouped by type and formatted into a natural-language context block prepended to the system prompt sent to Qwen via Alibaba Cloud Model Studio. The prompt explicitly instructs the model to treat RULE memories as binding behavioral instructions and all other types as context, a distinction that matters enormously in practice and that generic memory systems typically ignore.

Tech Stack

  • Backend: FastAPI, Python 3.11
  • Memory store: MongoDB (records + graph edges), Qdrant (384-dim embeddings), PostgreSQL with pgvector, Redis (retrieval cache)
  • ML: PyTorch, PyTorch Geometric (GraphSAGE), sentence-transformers (all-MiniLM-L6-v2)
  • LLM: Qwen via Alibaba Cloud Model Studio
  • Frontend: Streamlit
  • Infrastructure: Docker Compose, Alibaba Cloud ECS

Challenges

The user identity problem. The single most damaging bug in the project had nothing to do with machine learning. Generating a fresh random user_id per session meant the system was structurally incapable of accumulating memory regardless of how well every other component worked. Memory persistence requires persistent identity. Obvious in retrospect, invisible until you step back from the code and ask why the graph never grows.

Silent payload omissions. The memory write pipeline stored embeddings in Qdrant without including the raw text content in the payload. Retrieval returned correct memory IDs and similarity scores. The LLM received empty strings for every memory context. No error was raised anywhere in the pipeline. This class of bug, where the system functions correctly by every measurable metric but produces wrong behavior at the boundary between components, is the hardest to find and the most important to design against.

The cold start problem. A new user has no memory graph. A sparse graph produces weak GNN signal. Weak GNN signal means the system runs on flat vector retrieval, no better than a basic RAG system. The solution is a confidence gate that blends GNN output with flat retrieval proportionally to how much real evidence exists:

$$\alpha = \min\left(1.0,\ \frac{N_{\text{real}}}{N_{\text{threshold}}}\right), \quad N_{\text{threshold}} = 15$$

$$\text{final_score} = \alpha \cdot \text{gnn_score} + (1 - \alpha) \cdot \text{flat_score}$$

The GNN runs from turn one but its output is only trusted in proportion to the evidence behind it. The transition from cold start to full personalization is a smooth handoff rather than a hard threshold the user can feel.

Deployment. First-time deployment to a real server on Alibaba Cloud ECS surfaced issues that local Docker Compose development does not, SSH key configuration, security group rules, the difference between a named Docker service hostname and localhost, and the realization that every port exposed in docker-compose.yml also needs to be opened in the cloud firewall independently. None of these are hard problems individually. All of them together, for the first time, on a deadline, are a meaningful challenge.

Accomplishments to be proud off

Applying GNNs to a problem where they had not been applied before. A thorough search of prior work in academic papers, open source projects, competition entries, found no established use of graph neural networks for personal AI agent memory retrieval. GraphSAGE was not the obvious or default choice for this problem. It was a deliberate architectural hypothesis, built and validated from scratch.

A working GraphSAGE model integrated into a live inference pipeline. The model is not a notebook experiment. It runs inside a containerized FastAPI application on every retrieval call, produces enriched 128-dimensional node embeddings, and outputs relevance scores and cluster logits that feed directly into the hybrid scoring formula. The dual supervision loss, binary cross-entropy on relevance labels combined with categorical cross-entropy on memory type, prevents the shared embedding from collapsing to shortcuts that ignore the graph's semantic structure.

A multi-database architecture where every store serves a distinct access pattern. MongoDB holds canonical memory records and the graph edge collection. Qdrant handles 384-dimensional vector search filtered at query time by user. Redis caches retrieval results with a short TTL that respects the mutability of memory state. Each store was chosen for a specific reason, and no two stores are redundant.

A retrieval pipeline that combines four independent signals into a single learned score. Cosine similarity from Qdrant, GNN relevance from GraphSAGE, recency decay computed from turn age, and cluster confidence from the softmax head are blended using weights that reflect the relative trustworthiness of each signal. This is a meaningful departure from the standard RAG pattern of ranking by similarity alone.

A principled forgetting mechanism that archives rather than deletes. Memory archival is not a cleanup job. It is a core feature of the system, the mechanism that keeps the active memory set relevant as history accumulates, without destroying information that might matter later. The decay is gradual, the archival is reversible, and the forgetting schedule is tied to the session turn counter so it runs automatically without external orchestration.

What's Next for Mnemosyne

Mnemosyne as it stands is a working foundation. What comes next is about making it sharper, more autonomous, and ready for real multi-user deployment.

GNN training from real interaction data. The GraphSAGE model currently initializes with random weights and uses a heuristic relevance proxy for supervision. As Mnemosyne accumulates real interaction history, which memories were retrieved, which led to good responses, which were accessed repeatedly, those signals become genuine training labels. A background training worker that continuously fine-tunes the GNN on real access logs will make the relevance scoring increasingly accurate over time rather than static.

Multi-user scaling. The current architecture is correct for multi-user use but has only been validated with a single persistent user. Moving to real multi-user deployment requires per-user graph isolation at scale, more aggressive caching, and background workers for aging and forgetting that do not block the main request path.

Memory export and portability. A user's memory graph should belong to them. The next version will include an export mechanism, a structured JSON or graph format that lets a user take their accumulated memory out of the system, inspect it, or import it into another instance. Memory portability is a property that matters both ethically and practically.

The long-term vision is for Mnemosyne to become the memory layer that any AI agent can be built on, not a product, but infrastructure. The same way a database is not an application, Mnemosyne is not an assistant. It is what makes an assistant worth coming back to.

Built With

Share this project:

Updates