How to design a memory system for AI agents

Most agents break the moment a conversation spans more than a few turns because the context window is treated as the only state. To design AI agent memory system that survives real workloads, you need explicit tiers, controlled write paths, and retrieval that respects token budgets. This guide lays out an ordered path from prototype to production-grade memory.

1. Separate memory into three tiers

Working memory is the tokens currently in the prompt. Episodic memory is the raw, append-only log of what the agent and user did. Semantic memory is distilled facts and preferences extracted from episodes.

If you conflate them, you either truncate important history or drown the model in noise. A clean separation lets you apply different retention, storage, and cost policies per tier. When you design AI agent memory system, start by drawing these boundaries before writing any persistence code.

class MemoryTiers:
    def __init__(self):
        self.working = []   # list of dicts, last N messages in context
        self.episodic = []  # append-only event log, backed by durable store
        self.semantic = {}  # key-value facts, user_id -> extracted facts

Tradeoff: semantic memory requires extraction logic that can introduce drift. Keep the episodic log as the source of truth so you can rebuild semantics after a bad extraction run.

Why working memory is not a database

The context window is the fastest memory you have—zero lookup latency—but it is volatile and small. Treat it as a cache, not a store. Everything that matters must exist in episodic or semantic before the process restarts.

2. Define explicit write paths

Agents should not silently persist everything. Add a deterministic policy: after each agent step, call a should_store check, then write to episodic. Later, a background job promotes episodic entries to semantic.

Pitfall: writing raw LLM transcripts directly into semantic store multiplies embedding cost and pollutes retrieval with redundant phrasing. Store raw text in episodic; store only extracted facts in semantic.

def should_store(event: dict) -> bool:
    # never store heartbeat or tool ping events
    return event.get("type") not in ("heartbeat", "ping")

def write_episodic(log, event):
    if should_store(event):
        log.append({"ts": event["ts"], "data": event["data"]})

For fact extraction, use a small model call with a strict schema. Below is a minimal OpenAI-compatible call (works against any gateway that exposes /v1/chat/completions):

from openai import OpenAI
import json

client = OpenAI(base_url="https://api.n4n.ai/v1", api_key="sk-...")

def extract_facts(transcript: str) -> list[str]:
    resp = client.chat.completions.create(
        model="anthropic/claude-3-haiku",
        messages=[
            {"role": "system", "content": "Extract durable user facts as JSON list of strings."},
            {"role": "user", "content": transcript}
        ],
        response_format={"type": "json_object"}
    )
    return json.loads(resp.choices[0].message.content)["facts"]

Because n4n.ai provides automatic fallback when a provider is rate-limited or degraded, batch extraction jobs keep running even if the primary model throws 429s. That resilience matters when you run nightly summarization across thousands of sessions.

3. Pick storage backends per tier

Working memory stays in process memory; it is ephemeral and dies with the worker. Episodic needs an append-only store with time indexing. SQLite is enough for most single-node agents; use Kafka or a write-ahead log if you fan out across workers.

Semantic needs queryable vectors plus metadata. Use a vector DB (pgvector, Qdrant) but keep a relational table for permissions and TTL.

CREATE TABLE episodic (
    id INTEGER PRIMARY KEY,
    agent_id TEXT,
    ts INTEGER,
    payload TEXT
);
CREATE TABLE semantic_facts (
    id INTEGER PRIMARY KEY,
    agent_id TEXT,
    fact TEXT,
    embedding VECTOR(1536),
    updated_at INTEGER
);

Tradeoff: vector DBs add operational burden and a new failure mode. If your fact count is below 10k, a SQLite JSON column with lexical search and manual tagging often beats a vector index on latency and simplicity.

Redaction before write

Episodic logs are a liability surface. Strip credentials, PII, and raw tokens in the write_episodic step, not after. A simple regex pass is better than nothing:

import re
def redact(text: str) -> str:
    return re.sub(r"sk-[a-zA-Z0-9]{20,}", "[REDACTED]", text)

4. Retrieve before you construct the prompt

Never inject all memory. Run a retrieval step: embed the current user query, pull top-k semantic facts, and filter episodic by a sliding time window. The goal is to design AI agent memory system where the prompt contains only what changes the answer.

def retrieve(client, query, agent_id, k=5):
    q_emb = embed(client, query)
    facts = vector_search(agent_id, q_emb, k)
    recent = sqlite_select_recent(agent_id, limit=20)
    return format_for_prompt(facts, recent)

Pitfall: returning 20 chunks because “more context is safer” destroys token budget and degrades reasoning. Start with k=3 and measure task accuracy on a held-out set. Retrieval is a knob, not a floodgate.

Hybrid filtering

Pure vector search misses exact matches like order IDs. Add a metadata filter: WHERE agent_id = ? AND ts > ?. This cuts irrelevant hits and keeps cost predictable.

5. Run compression and summarization loops

Episodic logs grow unbounded. Schedule a job that summarizes older episodes into semantic facts. Use a stronger model for summarization, but keep it cheap with batching.

def summarize_batch(episodes):
    text = "\n".join(e["payload"] for e in episodes)
    return extract_facts(text)  # reuses earlier function

Set a retention policy: delete episodic older than 30 days unless flagged for audit. Semantic facts get a TTL and a re-validation prompt that re-checks against recent episodes.

Tradeoff: summarization loses nuance. If a user corrects a fact, the episodic log must win over the stale semantic entry. Implement a fact_version counter and let fresh episodes override.

6. Control tokens with cache hints and routing

Providers support cache control on prompt prefixes. If your gateway honors client routing directives and forwards provider cache-control hints, set them on static system prompts and long semantic blocks. This avoids re-paying for stable context on every turn.

curl https://api.n4n.ai/v1/chat/completions \
  -H "authorization: Bearer $KEY" \
  -H "x-cache-ttl: 3600" \
  -d '{"model":"openai/gpt-4o","messages":[{"role":"system","content":"You are a support agent. Stable policies: ..."}]}'

Per-token metering lets you attribute memory retrieval cost to the agent session. When a session blows the budget, you can trace it to oversized semantic pulls rather than guessing.

7. Common pitfalls and tradeoffs

• Over-retention: Storing every tool call creates privacy exposure and cost. Redact secrets before episodic write.
• Staleness: Semantic facts become wrong. Add a last_verified column and a weekly re-check loop.
• Latency: Retrieval adds a network round trip. Cache embeddings for repeated queries within a session.
• Drift: Extraction models hallucinate facts. Log confidence and require human review for high-impact facts like billing preferences.
• Lock-in: Custom vector schemas are hard to migrate. Keep the embedding model name in the row so you can re-embed later.

8. Minimal reference flow

1. Agent step produces an event.
2. should_store filters noise → write_episodic with redaction.
3. Background worker batches episodes → extract_facts → semantic upsert.
4. On new user query, retrieve builds constrained context with k=3–5.
5. Prompt sent with cache hints and per-token metering.
6. Summarization prunes episodic older than retention window.

That ordered path is how you design AI agent memory system that stays cheap, auditable, and correct under production load. Memory is not a feature you bolt on; it is the agent’s operating system. Build the tiers first, wire writes explicitly, and treat retrieval as a performance-critical query path.