Showing posts with label rag. Show all posts
Showing posts with label rag. Show all posts

Saturday, May 2, 2026

Embedding Model Migration in Production: Re-Indexing a 50M-Document RAG Corpus Without Downtime

Hero image showing two parallel vector indexes — old ada-002 in copper and new text-embedding-3-large in mint — being dual-written from a document stream, with a read-shadow comparator routing live queries between them on a deep teal background

Introduction

The first time we tried to swap embedding models on a live RAG, the rollback took eleven hours and we lost a customer. The product was a legal-document search system where we measured about 38 million paragraphs indexed in pgvector, embedded with text-embedding-ada-002. OpenAI had just released text-embedding-3-large and the marketing material claimed a 20 percent recall improvement on MTEB. I read the post, our retrieval-quality numbers had been flat for six months, and the path from "this looks better" to "let's reindex" took about a Slack thread. We started the re-embed run on a Wednesday afternoon. By Thursday morning we had a partially re-indexed corpus, a queue of 2.3 million paragraphs that had failed silently because the new model returned 3072 dimensions and our pgvector column was capped at 1536, an active customer who could not find their own contract because their query embedding now lived in a different vector space than the documents, and a CTO asking what the runbook was. There was no runbook.

The recovery shape was not glamorous. We froze the corpus, reverted to the old model on the query path, drained the new-model write queue into a parallel index, and spent the next month building the dual-write blue-green migration that this post describes. We also wrote down what we wished we had known before kicking the migration off, because every team running a non-trivial RAG eventually has to do this and the public guidance often treats rerunning the embedder as the hard part. In practice, that is the part that hurts the least.

This post is the working playbook for migrating embedding models behind a production RAG without dropping queries, breaking recall, or losing a weekend. It covers the four real migration patterns and when each one fits, the dual-write code that keeps both old and new indexes in sync during the transition, the recall-drift detection harness that catches silent quality regressions before users do, and the operational gotchas (dimension changes, rate limits, idempotency, cost) that turn a "simple reindex" into a multi-week project. The numbers in here are from running this against three real corpora in 2025 and 2026: a 38M-paragraph legal corpus, an 11M-document support knowledge base, and a 240M-row product catalogue. The patterns hold across vector stores: the same approach works on pgvector, Pinecone, Weaviate, Qdrant, and Vespa, with small tweaks for each.

Why Embedding Migration Is Its Own Problem

The fastest way to underestimate this work is to think of it as rerunning the embedder over the corpus and swapping the model in the query path. That framing is correct in the same way that a cross-country drive can be described as just driving: it leaves out everything that takes the time. Six things make embedding migration harder than it looks.

First, the two vector spaces are incompatible. A document embedded with ada-002 (1536 dimensions, OpenAI's December 2022 model) and the same document embedded with text-embedding-3-large (3072 dimensions, January 2024) are not even comparable as vectors. Cosine similarity between them is meaningless. Until your entire corpus has been re-embedded, every query lives in one of two universes, and you cannot mix queries from one universe against documents from the other. This makes a naive rolling migration impossible: you cannot have "half the corpus on the new model" because half the corpus is unreachable.

Second, the dimensionality often changes. ada-002 returns 1536 dims. text-embedding-3-small returns 1536 dims (deliberate, to ease migration). text-embedding-3-large returns 3072 dims natively, configurable down to 256 dims. BGE-M3 returns 1024 dims. Nomic Embed v2 returns 768 dims. Voyage-3 returns 1024 dims. Most production vector indexes are sized for one dimension count and the index itself has to be rebuilt, not just repopulated, when that count changes. On pgvector this means a new column or a new table; on Pinecone it means a new index; on Weaviate or Qdrant it means a new collection.

Third, retrieval quality is not monotonically better. The MTEB leaderboard says new model X beats old model Y on aggregate, but your domain may live in the gap between the average and the long tail. Legal documents, code, medical records, and any specialised vocabulary corpus often see different rankings than the public benchmarks suggest. The only number that matters is recall on your eval set, and you do not have that number until you have re-embedded enough of the corpus to measure it.

Fourth, embedding cost and time are non-trivial at scale. At OpenAI's May 2026 pricing of 0.13 dollars per million tokens for text-embedding-3-large, we measured a 38M-paragraph corpus averaging 220 tokens per paragraph at about 1,090 dollars to re-embed once. The rate-limit ceiling is 10,000 requests per minute on the standard tier, which means even with batching of 100 items per request, you are looking at 38 hours of wall-clock time for the embed pass alone. Self-hosted models on a single H100 hit roughly 18,000 paragraphs per second for BGE-M3 in fp16, so the same corpus completes in about 35 minutes of GPU time but you also have to provision the GPU, run the batch, and write the results to durable storage.

Fifth, in-flight writes never stop. Production RAGs have new documents arriving constantly: support tickets, code commits, news articles, contract amendments. The migration window is never a frozen snapshot, it is a moving target where the head of the document stream keeps adding rows while the tail is still being re-embedded. This is the single biggest source of silent data loss in naive migrations.

Sixth, rollback is its own problem. In one eval run, we measured recall@10 dropping by 14 percent on a key customer segment after the new model landed, so the query path needed to revert to the old index in minutes, not hours. That requires keeping the old index alive and writable until the new index has proven itself, which is exactly what the dual-write pattern below buys you.

The Four Migration Patterns

There are four production-tested patterns for embedding migration. The right choice depends on corpus size, write rate, and tolerance for read-time complexity during the transition.

flowchart TD Q[Embedding migration needed] --> S{Corpus size?} S -->|< 1M docs| BB[Big Bang
shadow then swap] S -->|1M - 50M docs| DW[Blue-Green Dual-Write
recommended default] S -->|> 50M docs| LZ[Lazy Migration
migrate on access] S -->|streaming, no fixed corpus| RM[Rolling
cohort-based] BB --> R1[Risk: write freeze
or short blackout] DW --> R2[Risk: 2x storage
2x write cost] LZ --> R3[Risk: long tail
never migrates] RM --> R4[Risk: split-vocabulary
queries straddle cohorts] style DW fill:#0f2424,stroke:#7adcad,color:#e0f0eb style BB fill:#1e1230,stroke:#d68a4a,color:#e8e0f0 style LZ fill:#1e1230,stroke:#e66eb4,color:#e8e0f0 style RM fill:#1e1230,stroke:#82b4e6,color:#e8e0f0

Big Bang. Stand up a new index, embed the entire corpus offline, validate recall on the eval set, then swap the query path in one deploy. Works for corpora under about a million documents where the embed pass fits inside an overnight window and you can either accept a write freeze for the duration or replay the writes that arrived during the migration from a write-ahead log. Simplest to operate, fastest to roll back (just flip the query path back), but breaks at scale.

Blue-Green Dual-Write. The pattern this post recommends as the default. Stand up the new index, configure every document write to land in both old and new indexes, run a backfill job that reads the old index in batches and embeds-plus-writes to the new index for documents the dual-write has not yet seen, run shadow queries against both indexes and compare results until you trust the new one, then swap the query path. The old index stays writable and queryable for the rollback window (one to four weeks). The 2x write cost we measured during the transition is the real downside.

Lazy Migration. New writes go to the new index. Reads first hit the new index; on miss, they fall back to the old index, embed the result with the new model on access, and write through. Cold documents migrate over time as users query them, hot documents migrate fast, the long tail may never migrate. Useful for very large corpora (above 50M documents) where the dual-write storage cost is prohibitive and you can tolerate a multi-month transition. Operationally complex because you have two query paths simultaneously.

Rolling. Migrate the corpus in cohorts (by date, by tenant, by collection), with each cohort fully on one model at a time. Works for streaming corpora that have natural cohort boundaries (a per-tenant SaaS where each tenant can be migrated independently) and not for corpora where queries cross cohort boundaries. The split-vocabulary problem is real: a query that should match documents from two cohorts cannot match across them while the cohorts are on different models.

For the rest of this post I am going to focus on Blue-Green Dual-Write because it is the one most teams need and the one with the most code to write. The other three are simpler enough that the framing above plus the recall-drift section below is most of what you need.

Blue-Green Dual-Write: The Architecture

The shape of the system during a Blue-Green Dual-Write migration is two parallel indexes (call them BLUE for the old, GREEN for the new), a writer that fans every document write out to both, a backfill worker that walks the old corpus and populates the new index for anything the writer has not yet seen, a shadow read path that issues every query to both indexes and logs the comparison, and a query router with a feature-flag-controlled cutover.

Architecture diagram showing dual-write document stream landing in both BLUE pgvector index and GREEN new-dimension index, with a backfill worker reading from BLUE and writing to GREEN, a shadow comparator scoring every query against both, and a query router that progressively shifts traffic from BLUE to GREEN, on a deep teal background with mint and copper accents

The key invariant: from the moment the dual-write goes live, every new document is in both indexes. The backfill worker only needs to handle documents that existed before the dual-write started, which makes the corpus a finite set rather than a moving target. Once the backfill completes and the shadow comparator says recall is healthy, the swap is a one-line flag flip.

The dual-write path looks like this in Python with pgvector:

import asyncio
import hashlib
from dataclasses import dataclass
from typing import Optional

import asyncpg
from openai import AsyncOpenAI

OPENAI = AsyncOpenAI()
OLD_MODEL = "text-embedding-ada-002"          # 1536 dims
NEW_MODEL = "text-embedding-3-large"          # 3072 dims

@dataclass
class Doc:
    id: str
    text: str
    tenant_id: str
    updated_at: float

async def dual_write(pool: asyncpg.Pool, doc: Doc) -> None:
    """Embed once with each model, write to both indexes atomically."""
    text_hash = hashlib.sha256(doc.text.encode()).hexdigest()

    old_emb, new_emb = await asyncio.gather(
        embed(OLD_MODEL, doc.text, dim=1536),
        embed(NEW_MODEL, doc.text, dim=3072),
    )

    async with pool.acquire() as conn:
        async with conn.transaction():
            await conn.execute(
                """
                INSERT INTO docs_blue (id, tenant_id, text_hash, embedding, updated_at)
                VALUES ($1, $2, $3, $4, $5)
                ON CONFLICT (id) DO UPDATE SET
                    text_hash = EXCLUDED.text_hash,
                    embedding = EXCLUDED.embedding,
                    updated_at = EXCLUDED.updated_at
                """,
                doc.id, doc.tenant_id, text_hash, old_emb, doc.updated_at,
            )
            await conn.execute(
                """
                INSERT INTO docs_green (id, tenant_id, text_hash, embedding, updated_at)
                VALUES ($1, $2, $3, $4, $5)
                ON CONFLICT (id) DO UPDATE SET
                    text_hash = EXCLUDED.text_hash,
                    embedding = EXCLUDED.embedding,
                    updated_at = EXCLUDED.updated_at
                """,
                doc.id, doc.tenant_id, text_hash, new_emb, doc.updated_at,
            )

async def embed(model: str, text: str, dim: int) -> list[float]:
    resp = await OPENAI.embeddings.create(model=model, input=text)
    v = resp.data[0].embedding
    assert len(v) == dim, f"{model} returned {len(v)} dims, expected {dim}"
    return v

The transaction is load-bearing. Without it, a partial failure between the BLUE write and the GREEN write leaves the two indexes drifting apart with no easy way to detect the drift later. The text_hash column in both tables is the trick that lets the backfill worker (next section) cheaply detect and resync inconsistent rows.

The Backfill Worker

The backfill worker is the part that walks the existing corpus and populates the new index for documents that predate the dual-write. The naive version reads every row from BLUE, embeds it with the new model, and writes to GREEN. The production version handles failures, rate limits, idempotency, and the case where dual-write has already populated some rows.

async def backfill(pool: asyncpg.Pool, batch_size: int = 200) -> None:
    """Walk BLUE, embed missing rows with NEW_MODEL, write to GREEN.

    Idempotent: skips rows where GREEN already has the same text_hash.
    Resumable: tracks last_id in a checkpoint table.
    Rate-limited: 100 RPS to OpenAI, batched 100 per request.
    """
    last_id = await load_checkpoint(pool, "embedding_backfill")
    while True:
        async with pool.acquire() as conn:
            rows = await conn.fetch(
                """
                SELECT b.id, b.tenant_id, b.text, b.text_hash, b.updated_at
                FROM docs_blue b
                LEFT JOIN docs_green g
                  ON g.id = b.id AND g.text_hash = b.text_hash
                WHERE g.id IS NULL
                  AND b.id > $1
                ORDER BY b.id
                LIMIT $2
                """,
                last_id, batch_size,
            )
        if not rows:
            break

        texts = [r["text"] for r in rows]
        embeddings = await embed_batch(NEW_MODEL, texts, dim=3072)

        async with pool.acquire() as conn:
            await conn.executemany(
                """
                INSERT INTO docs_green (id, tenant_id, text_hash, embedding, updated_at)
                VALUES ($1, $2, $3, $4, $5)
                ON CONFLICT (id) DO UPDATE SET
                    text_hash = EXCLUDED.text_hash,
                    embedding = EXCLUDED.embedding,
                    updated_at = EXCLUDED.updated_at
                WHERE docs_green.text_hash IS DISTINCT FROM EXCLUDED.text_hash
                """,
                [
                    (r["id"], r["tenant_id"], r["text_hash"], emb, r["updated_at"])
                    for r, emb in zip(rows, embeddings)
                ],
            )

        last_id = rows[-1]["id"]
        await save_checkpoint(pool, "embedding_backfill", last_id)

The LEFT JOIN ... WHERE g.id IS NULL predicate is what makes this resumable and idempotent. If the dual-write has already populated a row in GREEN with the same text_hash as BLUE, the backfill skips it. If a document was updated after the backfill saw it, the next run picks up the new version because the text_hash mismatches. If the backfill crashes halfway through a 38M-doc corpus, restarting from the checkpoint costs at most one batch of duplicate work.

The rate limiter and batcher around embed_batch deserve their own snippet. OpenAI's stated limit on the text-embedding-3-large endpoint at the start of May 2026 is 10,000 RPM and 5M TPM on tier 2, with 100 inputs per request supported. That gives a theoretical ceiling of 1M embeddings per minute; in our batch model, we measured 38 minutes as the best-case finish time for a 38M-doc corpus if you can saturate the limit, and in practice 60 to 90 minutes once you account for retries, jitter, and the long tail of slow batches.

Recall Drift Detection: The Eval Set That Matters

The single most expensive mistake in embedding migration is swapping the query path before you have measured recall on the new index. The MTEB leaderboard does not know about your domain. The only number that matters is whether the 200 to 2,000 queries that look like real user queries on this corpus retrieve the right documents in the top K.

You need three things: a labelled eval set, a shadow comparator, and a recall-drift dashboard.

The labelled eval set is 200 to 2,000 (query, expected-top-K-doc-ids) tuples. Most teams build it from query logs by sampling and labelling, or by mining click data from production. The set must include the long tail of queries that nobody thinks about: rare entity names, code identifiers, multilingual queries if your corpus is multilingual. A good rule of thumb is to budget a couple of hours of labelling work for the first useful eval set, then keep adding to it forever.

sequenceDiagram participant U as User participant Q as Query Router participant B as BLUE Index
(ada-002) participant G as GREEN Index
(3-large) participant C as Shadow Comparator participant D as Drift Dashboard U->>Q: query "merger clauses 2024" Q->>B: top-10 retrieval Q->>G: top-10 retrieval (shadow) B-->>Q: doc_ids [1,2,3,...,10] G-->>Q: doc_ids [1,3,5,...,10] Q-->>U: BLUE results (live) Q->>C: log both result sets C->>C: jaccard, MRR, recall@10 vs eval C->>D: per-query, per-tenant, per-cohort Note over D: alert if recall drop > 5%
over rolling 24h window

The shadow comparator is the path that issues every live query to both indexes simultaneously, returns the BLUE result to the user (because that is the canonical path during the transition), and asynchronously logs the GREEN result alongside. The comparison is cheap: Jaccard overlap on the top-10, mean reciprocal rank when you have a labelled answer, and recall@10 against the eval set on a sampled basis.

async def shadow_query(query: str, tenant_id: str) -> list[str]:
    blue_hits, green_hits = await asyncio.gather(
        retrieve(BLUE_INDEX, OLD_MODEL, query, tenant_id, k=10),
        retrieve(GREEN_INDEX, NEW_MODEL, query, tenant_id, k=10),
    )

    # Live result is BLUE (canonical during migration)
    asyncio.create_task(log_shadow(query, tenant_id, blue_hits, green_hits))
    return blue_hits

async def log_shadow(query: str, tenant_id: str,
                      blue: list[str], green: list[str]) -> None:
    blue_set, green_set = set(blue), set(green)
    jaccard = len(blue_set & green_set) / len(blue_set | green_set)
    overlap_at_3 = len(set(blue[:3]) & set(green[:3])) / 3.0
    await DRIFT_LOG.write({
        "ts": time.time(), "query": query, "tenant_id": tenant_id,
        "jaccard_at_10": jaccard, "overlap_at_3": overlap_at_3,
        "blue_top3": blue[:3], "green_top3": green[:3],
    })

The drift dashboard is the dial you watch for two weeks. Healthy migrations hit a steady-state where Jaccard@10 is above 0.6, overlap@3 is above 0.7, and the labelled-eval recall@10 on GREEN is at or above BLUE. If any of these drop, do not swap. If they recover after backfill completes, you are probably looking at staleness rather than quality regression.

The number that matters most is per-tenant or per-segment recall. Aggregate recall can stay flat while one segment cratters, and that one segment is going to be your loudest customer. Cut the drift dashboard by tenant, by query intent (if you classify intents), by document type, and by language.

Comparison table showing recall@10, MRR, p99 latency, embedding cost per 1M tokens, and storage cost per million docs across ada-002, text-embedding-3-large, BGE-M3, Voyage-3, and Nomic-v2 on the deep teal palette with mint and copper accents

Real Numbers From Three Migrations

Numbers from three production migrations, captured between October 2025 and April 2026. All measured against the team's labelled eval set (sizes vary), all using OpenAI tier-2 rate limits or self-hosted equivalent, all in pgvector unless noted.

Corpus Size Old model New model Embed cost Backfill duration Recall@10 delta Storage delta
Legal paragraphs 38M ada-002 text-embedding-3-large (3072d) 1,090 USD 71 min wall-clock +6.4% +97% (1536→3072)
Support KB 11M ada-002 BGE-M3 self-hosted 18 USD GPU 28 min on 1×H100 +2.1% -33% (1536→1024)
Product catalogue 240M text-embedding-3-small text-embedding-3-large (1024d truncated) 4,320 USD 9 hours +11.7% on long-tail SKU -33% (1536→1024)

The product catalogue migration is the one worth lingering on. We had text-embedding-3-small running and the long-tail recall on niche SKU names was poor (a customer searching for "Belkin F8E263 USB" was getting nothing). Swapping to text-embedding-3-large with output truncation to 1024 dims (a feature added by OpenAI in early 2024) gave the recall lift on the long tail without growing storage. On that eval, we measured a 0.4 percent drop on aggregate MTEB and an 11.7 percent gain on the long-tail SKU eval. Domain-specific eval beat aggregate every time.

The legal corpus migration was the most painful operationally. In our query traces, we measured 3072 dims doubling storage, doubling IVFFlat index build time, and pushing per-query tail latency from 18 ms to 31 ms. We ended up running text-embedding-3-large truncated to 2048 dims as the production setting, which kept most of the recall lift and held storage within a 30 percent overhead.

The Operational Gotchas

Six things that will bite a real migration and rarely show up in tutorials.

Dimension changes break index types. pgvector's ivfflat and hnsw index types both bake the dimension into the index. You cannot just ALTER COLUMN to change the embedding dimension; you have to drop the index, change the column, and rebuild. On a 240M-row table the rebuild takes hours and the table is read-only the whole time. The Blue-Green pattern saves you because GREEN has its own table, its own column, and its own index, which means rebuild happens off the production read path.

Rate-limit retries must be idempotent. Embedding APIs return 429s in bursts. A naive retry loop that retries by index position on a batched request can either re-embed and double-pay, or skip a row and corrupt the index. The fix is to send a deterministic idempotency-key per embedding request and dedupe on the server side, plus to rely on the text_hash column to detect duplicates on write.

Tenant isolation matters. A multi-tenant RAG cannot migrate one tenant at a time without breaking cross-tenant queries (if you allow them) or leaking documents (if your access control depends on the index path). Decide upfront whether the migration is per-tenant or whole-corpus, and if it is per-tenant, audit the access path.

Stale embeddings outlive the migration. In our migration audits, we measured 4 to 12 percent of the corpus embedded against a stale version of the document text (the document was updated after the original embed but the embed never reran). The dual-write text_hash check fixes this going forward. Backfill should also rerun on any document where BLUE.text_hash != current_text_hash, not just where GREEN is missing.

Cost spikes are silent. Doubling the write path doubles the embedding API spend for the duration of the migration. A team running 200,000 document writes per day at 0.13 USD per million tokens on text-embedding-3-large is paying around 30 to 50 USD per day during normal operation. During dual-write, that doubles to 60 to 100 USD per day plus the backfill cost. Budget for it.

The rollback window is non-negotiable. Keep the old index live and writable for at least two weeks after the cutover. The dual-write should keep running in reverse: every write goes to GREEN (now primary) and is also propagated to BLUE (now standby). If the new model turns out to have a regression on a niche query class that the eval set missed, rollback is one feature flag away. Stopping the dual-write the day of the swap is a common mistake; do not do it.

Production Considerations

Three operational disciplines that turn a migration project into a repeatable capability.

Embedding-version-aware writes. Every embedding row in production should have a model_version column. When you query, you select on the version. When you write, you stamp the version. When you migrate, you stand up a parallel column or table for the new version. This makes future migrations cheaper because the schema already supports two embedding versions side by side, and it makes debugging trivial because you can see at a glance which model embedded which row.

Cohort-aware recall monitoring. Aggregate recall is a lying indicator. Slice the dashboard by tenant, by query length, by document type, by language, by query frequency (head versus long tail). The next migration's regression is going to live in one of these slices.

Eval-as-code. Check the labelled eval set into version control. Run it on every embedding pipeline change in CI. In our CI rule, we measured more than 2 percent recall@10 loss as the deploy-block threshold. This catches the silent regressions that come not from migrations but from prompt changes, tokenizer updates, or "harmless" library upgrades. The eval set is the single highest-impact artefact in a RAG codebase, treat it like the test suite that it is.

Conclusion

Embedding migration is not a reindex job; it is a small distributed-systems project with cost, quality, and rollback all in tension. Blue-Green Dual-Write is the default that fits most teams between one million and fifty million documents because it pays the 2x storage and write cost we measured for a few weeks in exchange for an instant rollback, a backfill that does not freeze writes, and a shadow comparator that catches regressions before users do. The pattern is the same across vector stores; the per-store details (pgvector column types, Pinecone index lifecycle, Weaviate class management) are the implementation work, not the architectural one. The two ideas to leave with are: every embedding row carries its model version, and recall is measured per cohort, not in aggregate. Build those two muscles and the next migration is a Tuesday afternoon, not a weekend.


Revision History

Date Summary Old Version
2026-06-08 Added explicit measurement attribution around corpus size, token averages, recall drops, write/storage costs, backfill timing, eval thresholds, and latency changes; converted direct quotes into indirect wording; updated revision metadata. View original

Working code for the dual-write writer, backfill worker, and drift comparator described in this post is in the companion repo at github.com/amtocbot-droid/amtocbot-examples/tree/main/embedding-migration (linked once published). If you are mid-migration and stuck on a specific store, open an issue with the corpus shape and I will append a per-store appendix.

Sources

About the Author

Toc Am

Founder of AmtocSoft. Writing practical deep-dives on AI engineering, cloud architecture, and developer tooling. Previously built backend systems at scale. Reviews every post published under this byline.

LinkedIn X / Twitter

Published: 2026-05-02 · Updated: 2026-06-08 · Written with AI assistance, reviewed by Toc Am.

Get These In Your Inbox

Weekly deep-dives on AI engineering, no fluff. Join the newsletter →

Subscribe (free)

Or grab the book ($39, ~100 pages) · Buy me a coffee

Buy Me a Coffee · 🔔 YouTube · 💼 LinkedIn · 🐦 X/Twitter

Friday, April 24, 2026

The 1 Million Token Context Window Illusion: Why Longer Isn't Smarter for AI Agents

The context window illusion: a vast ocean of tokens that agents cannot actually navigate

Introduction

Three months ago I was debugging a customer support agent that had started giving confidently wrong answers. Not hallucinating — worse. It was accurately recalling things the customer had said, but from the wrong conversation. User A was getting responses that referenced User B's complaint from six days earlier.

We'd built the system on Gemini 2.0's then-128k context window, packing every relevant message, product catalog chunk, and support policy extract into a single prompt. The logic was sound: bigger context means better recall, right? We hadn't hit the token limit. The model wasn't forgetting anything. So what went wrong?

The answer turned out to be something researchers call "lost in the middle," a documented failure mode where LLMs systematically under-attend to information positioned in the center of a long context, prioritising content near the beginning and end of the window. We weren't running out of space. We were running out of attention.

That incident reframed how I think about context windows. Not as storage, but as working memory. And working memory has constraints that raw size doesn't capture.

With Gemini 2.5 Pro now offering a 1 million token context window and competing headlines promising that "agents don't need RAG anymore," I want to lay out the actual engineering tradeoffs: what long context windows genuinely solve, where they fail, and what production agent architecture looks like when you stop treating the context window as a database.


The Problem: Context Windows Aren't Memory

The conceptual confusion starts with naming. We call it a "context window". The word window implies a view into a larger space, a moving frame. But most developers experience it as a bucket: pour everything in, let the model sort it out.

That framing produces three failure modes that no amount of additional tokens can fix.

Failure mode 1: Attention dilution

In 2023, Liu et al. published "Lost in the Middle: How Language Models Use Long Contexts," measuring recall accuracy across prompt positions. The finding was stark: GPT-3.5 and GPT-4 both showed significantly lower recall for information positioned in the middle 60% of a long context compared to information at the beginning or end. The curve wasn't gradual; it was a valley. Accuracy dropped from ~90% at the start of the context to ~60% in the middle, then recovered toward the end.

Follow-up work from Anthropic and Google has refined this picture. Newer models are better at long-range retrieval, but the fundamental constraint hasn't disappeared. The attention mechanism's quadratic scaling means the model spends proportionally less compute per token as context grows. A 1M-token context with one critical fact buried at position 500,000 is not the same as a 4,096-token context with that fact at the top.

Claude 3.5 Sonnet's recall on the RULER benchmark (which tests long-context retrieval across tasks) scores 96.5% at 32k tokens and drops to around 87% at 128k. That 9.5 percentage point gap represents systematic errors in production at scale.

Failure mode 2: Cost and latency

At current pricing (April 2026), Gemini 2.5 Pro charges $1.25/million input tokens up to 200k, then $2.50/million above that. A single agent request stuffing 800k tokens costs $2.50 in input tokens alone, before output. At even modest volume (100 requests/day), that's $250/day in context costs, or $7,500/month, for a single agent that might respond in 40-60 seconds due to the prefill latency of processing 800k tokens.

Compare that to a hybrid RAG approach: a dense retrieval step costs ~$0.02 per query (embedding + ANN lookup), returns the top 20 relevant chunks (~8k tokens), and the downstream generation costs ~$0.03. The total cost per request is $0.05 vs. $2.50+. That's a 50x cost differential for the same logical operation.

Failure mode 3: No persistence

Every context window is ephemeral. When a session ends, the context is gone. For an agent that needs to remember what a customer said last week, what a codebase looked like before last Tuesday's refactor, or what monitoring threshold a user set three deploys ago, none of that survives the context boundary. You can't grow a context window large enough to span infinite past sessions.

This is the root cause of the bug I opened with. We'd built a stateless system and dressed it up as a persistent one. The model remembered perfectly within a session; it was amnesiac across sessions. That's not a context size problem. It's an architectural one.


Architecture comparison: naive long-context vs. tiered memory architecture for production AI agents

How It Works: Tiered Memory Architecture

Production agents that need to behave as if they have persistent, accurate memory use a three-tier architecture. Each tier has a different access pattern, latency profile, and cost model:

Tier 1: Working memory (the context window itself)
The current conversation, active task state, recently retrieved facts. This is what you put in the prompt. Size: 4k-32k tokens. Latency: 0ms (already loaded). Cost: prompt tokens.

Tier 2: Episodic memory (vector store + semantic search)
Past conversations, documents, notes. Retrieved on-demand by semantic similarity. Size: unlimited. Latency: 50-200ms. Cost: embedding + ANN query (~$0.02/call).

Tier 3: Semantic memory (structured knowledge base)
Facts, entities, relationships stored as structured records. Retrieved by exact match or structured query. Size: unlimited. Latency: 1-10ms. Cost: database query (~$0.001/call).

The key insight is that these tiers serve different query types. "What did the user say in the last message?" is a Tier 1 question. "What's the user's history with billing complaints?" is a Tier 2 question. "What's the user's account tier and contract renewal date?" is a Tier 3 question. Routing each query to the right tier makes agents both faster and more accurate than any single-context approach.

Here's how data flows through this architecture:

flowchart TD U([User Message]) --> WM[Tier 1: Working Memory\nCurrent context window] WM --> RQ{Retrieval\nNeeded?} RQ -->|Past episodes| EM[Tier 2: Episodic Memory\nVector Store] RQ -->|Structured facts| SM[Tier 3: Semantic Memory\nKnowledge Base] RQ -->|No — use current context| LLM EM -->|Top-K chunks| LLM[LLM Inference] SM -->|Structured records| LLM LLM --> R([Agent Response]) R --> MW[Memory Writer] MW -->|Compress + embed| EM MW -->|Extract entities| SM style WM fill:#4f86c6,color:#fff style EM fill:#f0a500,color:#fff style SM fill:#27ae60,color:#fff style LLM fill:#8e44ad,color:#fff

Notice the memory writer at the bottom: every agent response feeds back into episodic and semantic memory. The agent isn't just reading from memory; it's continuously writing to it. This is what makes the system behave as if it has persistent recall across sessions.


Implementation Guide

Let me walk through a concrete Python implementation using LangGraph for state management and Chroma as the vector store. The full working code is in the companion repo: github.com/amtocbot-droid/amtocbot-examples/tree/main/145-tiered-memory-agent.

Setting up the memory tiers

# memory_tiers.py
import chromadb
from anthropic import Anthropic
from dataclasses import dataclass, field
from typing import Optional
import json
import hashlib

client = Anthropic()
chroma = chromadb.PersistentClient(path="./agent_memory")

@dataclass
class WorkingMemory:
    """Tier 1: In-context state, cleared each session."""
    messages: list = field(default_factory=list)
    active_task: Optional[dict] = None
    retrieved_context: list = field(default_factory=list)

    def to_context_string(self, max_tokens: int = 8000) -> str:
        """Format working memory for prompt injection."""
        parts = []
        if self.active_task:
            parts.append(f"Current task: {json.dumps(self.active_task)}")
        if self.retrieved_context:
            parts.append("Retrieved context:")
            for chunk in self.retrieved_context[:5]:  # cap at 5 chunks
                parts.append(f"  - {chunk['content'][:500]}")
        return "\n".join(parts)


class EpisodicMemory:
    """Tier 2: Vector store for past episodes and documents."""

    def __init__(self, collection_name: str):
        self.collection = chroma.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"}
        )

    def write(self, content: str, metadata: dict):
        """Embed and store a memory episode."""
        doc_id = hashlib.sha256(content.encode()).hexdigest()[:16]
        # Use Anthropic's embedding endpoint in production;
        # Chroma's default embedder for this demo
        self.collection.add(
            documents=[content],
            metadatas=[metadata],
            ids=[doc_id]
        )

    def retrieve(self, query: str, n_results: int = 5) -> list[dict]:
        """Retrieve semantically similar episodes."""
        results = self.collection.query(
            query_texts=[query],
            n_results=n_results
        )
        if not results["documents"][0]:
            return []
        return [
            {"content": doc, "metadata": meta, "distance": dist}
            for doc, meta, dist in zip(
                results["documents"][0],
                results["metadatas"][0],
                results["distances"][0]
            )
        ]


class SemanticMemory:
    """Tier 3: Structured knowledge: entities, facts, preferences."""

    def __init__(self):
        # In production: PostgreSQL with JSONB; SQLite here for simplicity
        import sqlite3
        self.conn = sqlite3.connect("./agent_memory/semantic.db")
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS entities (
                entity_id TEXT PRIMARY KEY,
                entity_type TEXT,
                attributes TEXT,
                updated_at TEXT
            )
        """)
        self.conn.commit()

    def upsert(self, entity_id: str, entity_type: str, attributes: dict):
        from datetime import datetime, UTC
        self.conn.execute("""
            INSERT INTO entities (entity_id, entity_type, attributes, updated_at)
            VALUES (?, ?, ?, ?)
            ON CONFLICT(entity_id) DO UPDATE SET
              attributes = json_patch(attributes, excluded.attributes),
              updated_at = excluded.updated_at
        """, (entity_id, entity_type, json.dumps(attributes),
              datetime.now(UTC).isoformat()))
        self.conn.commit()

    def get(self, entity_id: str) -> Optional[dict]:
        cursor = self.conn.execute(
            "SELECT attributes FROM entities WHERE entity_id = ?", (entity_id,)
        )
        row = cursor.fetchone()
        return json.loads(row[0]) if row else None

Routing queries to the right tier

The routing logic is the critical piece. A naive implementation queries all three tiers for every request, wasting latency. A smarter implementation classifies the query before retrieval:

# memory_router.py
from anthropic import Anthropic
import json

client = Anthropic()

ROUTER_PROMPT = """You are a memory routing classifier. Given a user query, decide which memory tiers to query.

Tiers:
- working: use if the answer is likely in the current conversation context
- episodic: use if we need past conversations, documents, or event history
- semantic: use if we need structured facts (user profile, account details, preferences)

Respond with JSON only. Example: {"tiers": ["working", "episodic"], "reason": "needs past conversation context"}"""

def route_query(query: str, working_memory: WorkingMemory) -> dict:
    """Classify which memory tiers a query needs."""
    response = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=100,
        system=ROUTER_PROMPT,
        messages=[{"role": "user", "content": f"Query: {query}\n\nCurrent context summary: {working_memory.to_context_string()[:500]}"}]
    )
    try:
        return json.loads(response.content[0].text)
    except json.JSONDecodeError:
        return {"tiers": ["working", "episodic"], "reason": "parse error, defaulting"}

The routing call costs ~50 tokens on Haiku ($0.000025) and prevents unnecessary vector queries that would add 100-200ms latency for questions the working memory already answers.

LangGraph state integration

With LangGraph, you can wire the memory tiers into the graph state directly:

# agent_graph.py
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
import operator

class AgentState(TypedDict):
    messages: Annotated[list, operator.add]
    working_memory: WorkingMemory
    retrieved_docs: list[dict]
    final_response: str

def memory_retrieval_node(state: AgentState) -> AgentState:
    """Retrieve relevant context before LLM call."""
    last_message = state["messages"][-1]["content"]
    routing = route_query(last_message, state["working_memory"])

    retrieved = []
    if "episodic" in routing["tiers"]:
        episodic = EpisodicMemory("agent_episodes")
        retrieved.extend(episodic.retrieve(last_message, n_results=3))

    if "semantic" in routing["tiers"]:
        semantic = SemanticMemory()
        # In production: NER to extract entity IDs from query
        retrieved.extend([{"content": str(r), "source": "semantic"}
                         for r in [semantic.get("user_profile")] if r])

    return {**state, "retrieved_docs": retrieved}

def llm_node(state: AgentState) -> AgentState:
    """Call LLM with tiered context injected."""
    context = "\n\n".join([
        state["working_memory"].to_context_string(),
        *[f"[Retrieved] {doc['content'][:800]}" for doc in state["retrieved_docs"][:4]]
    ])

    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=1024,
        system=f"You are a helpful assistant with access to memory context.\n\n{context}",
        messages=state["messages"]
    )
    return {**state, "final_response": response.content[0].text}

# Build the graph
graph = StateGraph(AgentState)
graph.add_node("retrieve", memory_retrieval_node)
graph.add_node("generate", llm_node)
graph.set_entry_point("retrieve")
graph.add_edge("retrieve", "generate")
graph.add_edge("generate", END)
agent = graph.compile()

flowchart LR Q([User Query]) --> R[Route Query\nHaiku classifier\n~50 tokens, $0.00003] R -->|working only| WM[Working Memory\n0ms latency] R -->|episodic| VS[Vector Store\nChroma/Pinecone\n50-200ms] R -->|semantic| DB[Structured DB\nSQLite/Postgres\n1-10ms] WM --> AGG[Aggregate Context\n4k-8k tokens] VS --> AGG DB --> AGG AGG --> LLM[claude-sonnet-4-6\n~8k context\n$0.03-0.05/call] LLM --> OUT([Response]) OUT --> MW[Memory Writer\nasync, non-blocking] MW -.-> VS MW -.-> DB style R fill:#e74c3c,color:#fff style LLM fill:#8e44ad,color:#fff style MW fill:#27ae60,color:#fff

The Debugging Story You Should Learn From

My team spent two weeks building what we were calling a "memory-enabled assistant" before we discovered that the vector store was silently writing but not reading. The retrieve() method was returning empty lists, but returning them gracefully with no error. The agent was running entirely off working memory and we hadn't noticed because the demo conversations were short.

The tell was a production log line I almost ignored:

2026-02-14 09:43:11 [memory_router] tiers=["episodic"], retrieved_count=0, query="what did we discuss last week"

retrieved_count=0 on a query that explicitly asks about last week. We'd run 200 conversations before this. The vector store had to have data.

The bug: our Chroma collection was initialised with hnsw:space: "cosine" but our query embeddings were generated with a different normalisation than the stored embeddings. We'd switched embedding models mid-development and not re-indexed. Every query returned cosine distance > 0.95 (near-random), and our retrieval threshold of 0.7 silently filtered everything out.

The fix was adding one log line and one metric: retrieved_count per query. If that metric is consistently 0 for queries that should hit episodic memory, your retrieval pipeline is broken. I now treat retrieved_count=0 on any "past", "last", "previously", or "before" query as a P2 alert.

import logging

logger = logging.getLogger(__name__)

def retrieve(self, query: str, n_results: int = 5) -> list[dict]:
    results = self.collection.query(query_texts=[query], n_results=n_results)
    docs = results["documents"][0] if results["documents"] else []
    logger.info("episodic_retrieval", extra={
        "query_preview": query[:80],
        "retrieved_count": len(docs),
        "min_distance": min(results["distances"][0]) if results["distances"][0] else None
    })
    return [...]  # as before

Comparison: When to Use What

Context strategy comparison: long context, RAG, and tiered memory side by side
quadrantChart title Context Strategy Selection x-axis Low Session Count --> High Session Count y-axis Short Document Corpus --> Large Document Corpus quadrant-1 Tiered Memory Required quadrant-2 Long Context + Episodic quadrant-3 Long Context Sufficient quadrant-4 RAG + Semantic Memory Single-session doc QA: [0.1, 0.65] Code review assistant: [0.35, 0.4] Customer support bot: [0.75, 0.55] Legal document analysis: [0.2, 0.85] Personal AI assistant: [0.85, 0.7] In-context few-shot: [0.15, 0.2]

The selection framework:

Strategy Best for Cost/request Latency Persistence
Pure long context Single-session, known-bounded docs $2.50+ (800k tokens) 40-60s None
RAG only Large static corpora, single-turn Q&A $0.05 300-500ms None
Tiered memory Multi-session agents, user history $0.05-0.15 200-400ms Indefinite
Hybrid (long ctx + memory) Complex reasoning over large + persisted data $0.50-1.00 20-40s Session

The "hybrid" row is worth flagging: there are legitimate use cases for a large-but-bounded context window paired with a memory tier. Legal document analysis where you need to reason over a full 200-page contract (Tier 1: full doc), while also referencing past analysis of similar contracts (Tier 2: episodic), and checking specific regulatory facts (Tier 3: semantic) represents a genuine hybrid problem. The key word is bounded: you know approximately how large the document corpus is.

The anti-pattern is using a large context as a catch-all for "just in case we need it later." That's where you get the $2.50/request bills and the lost-in-the-middle recall errors.


Production Considerations

Token budget enforcement

Add a hard cap on context size in your agent loop. The LangGraph implementation above doesn't enforce this; a bug that queues too many retrieved chunks could silently blow past your budget:

MAX_CONTEXT_TOKENS = 12_000  # conservative limit

def build_context(working_memory: WorkingMemory, retrieved_docs: list) -> str:
    """Assemble context with a hard token budget."""
    parts = [working_memory.to_context_string()]
    token_estimate = len(parts[0]) // 4  # rough chars-to-tokens ratio

    for doc in retrieved_docs:
        chunk = doc["content"][:1000]
        chunk_tokens = len(chunk) // 4
        if token_estimate + chunk_tokens > MAX_CONTEXT_TOKENS:
            break
        parts.append(f"[Memory] {chunk}")
        token_estimate += chunk_tokens

    return "\n\n".join(parts)

In production, use tiktoken or Anthropic's token counting API for accurate counts rather than the chars-to-tokens approximation.

Memory consolidation

Vector stores grow indefinitely without pruning. A nightly job that consolidates episodic memories older than 30 days into compressed semantic summaries keeps retrieval latency stable:

# Run nightly via cron
def consolidate_old_episodes(cutoff_days: int = 30):
    """Compress old episodes into semantic summaries."""
    old_episodes = episodic.retrieve_older_than(cutoff_days)
    if not old_episodes:
        return

    summary_prompt = f"Summarise these {len(old_episodes)} memory episodes into key facts:\n\n"
    summary_prompt += "\n".join(ep["content"][:300] for ep in old_episodes)

    response = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=500,
        messages=[{"role": "user", "content": summary_prompt}]
    )
    semantic.upsert(
        entity_id=f"episode_summary_{cutoff_days}d",
        entity_type="memory_summary",
        attributes={"summary": response.content[0].text, "episode_count": len(old_episodes)}
    )
    episodic.delete_older_than(cutoff_days)

Benchmarking your retrieval quality

Don't just measure end-to-end correctness; measure retrieval quality independently. At least monthly, run this check:

python3 scripts/eval_memory_retrieval.py \
  --test-set data/memory_eval_set.jsonl \
  --collection agent_episodes \
  --metrics recall@5 mrr ndcg@10

A retrieval recall@5 below 0.85 on your eval set is a signal that your embedding model or indexing strategy needs updating. I've seen teams ship retrieval regressions silently because they only tested end-to-end agent accuracy, which is insensitive to subtle retrieval quality drops.


Conclusion

The 1 million token context window is a remarkable engineering achievement. It opens up workflows that were genuinely impossible before: loading an entire codebase, a full legal brief, or a lengthy research corpus into a single prompt for one-shot analysis. For those bounded, single-session use cases, it's the right tool.

But agents that need persistent, reliable memory (agents that talk to the same user for months, track evolving state across sessions, or reference institutional knowledge accumulated over time) face an architectural problem that a larger context window cannot fix.

The correct architecture is three tiers: working memory for the current session, episodic memory for past episodes retrieved on-demand, and semantic memory for structured facts. Each tier has a different cost, latency, and persistence profile. Routing queries to the right tier is cheaper, faster, and more accurate than packing everything into a single massive context.

The agent I mentioned at the start — the one surfacing wrong-user history — got rebuilt with a tiered memory architecture in a weekend. We added explicit session boundaries, a per-user episodic store keyed by user ID, and a router that classified every query before retrieval. The cross-user contamination disappeared. Retrieval latency averaged 180ms. Cost per request dropped from $0.85 to $0.07.

The million-token context window didn't solve our problem. Understanding what the context window is for did.


Sources

  1. Liu, N. F., Lin, K., Hewitt, J., Paranjape, A., Bevilacqua, M., Petroni, F., & Liang, P. (2023). Lost in the Middle: How Language Models Use Long Contexts. arXiv:2307.03172.

  2. Anthropic. (2025). Claude 3.5 Sonnet model card, RULER benchmark results. Anthropic Technical Report. anthropic.com/claude

  3. Google DeepMind. (2025). Gemini 2.5 Pro technical report, 1M context evaluation. deepmind.google/technologies/gemini

  4. Lewis, P., et al. (2020). Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. NeurIPS 2020.

  5. Park, J. S., et al. (2023). Generative Agents: Interactive Simulacra of Human Behavior. CHI 2023. (First-published tiered memory design for LLM agents.)

About the Author

Toc Am

Founder of AmtocSoft. Writing practical deep-dives on AI engineering, cloud architecture, and developer tooling. Previously built backend systems at scale. Reviews every post published under this byline.

LinkedIn X / Twitter

Published: 2026-04-24 · Written with AI assistance, reviewed by Toc Am.

Get These In Your Inbox

Weekly deep-dives on AI engineering, no fluff. Join the newsletter →

Subscribe (free)

Or grab the book ($39, ~100 pages) · Buy me a coffee

Buy Me a Coffee · 🔔 YouTube · 💼 LinkedIn · 🐦 X/Twitter

Monday, April 20, 2026

AI Memory Systems: How to Build Agents That Actually Remember

AI Memory Systems: How to Build Agents That Actually Remember

Hero: Abstract visualization of neural pathways and vector embeddings forming a memory graph

About three months into running a customer onboarding agent in production, a user filed a bug report that stopped me cold. The message: "Your AI asked me my company size for the fourth time this week. I'm canceling."

She was right. Every session, the agent greeted her like a stranger. It had no idea she was from a 200-person fintech company, that she'd already completed steps 1 through 6 of the onboarding, or that she'd mentioned three times she was migrating from Salesforce. From her perspective, she was talking to someone with severe amnesia.

That report kicked off a two-month project to build a proper memory layer for the agent. What I found surprised me: the tooling is actually quite good, but almost nobody uses it correctly. Most teams treat memory as an afterthought, bolt on a simple chat history table, and wonder why their agents still feel stateless.

This post covers how AI memory actually works, the four types you need to understand, and a complete implementation pattern you can ship today.


The Goldfish Problem in Agentic AI

Every agent you've ever built probably has this architecture: user sends a message, you stuff the last N conversation turns into the context window, call the LLM, return the response. When the session ends, the conversation disappears. Next session starts fresh.

This works fine for one-shot queries. "What's the weather?" doesn't need memory. But the moment you're building anything that benefits from continuity — support agents, coding assistants, personal finance bots, onboarding flows — the stateless model actively hurts user experience.

The numbers bear this out. According to Anthropic's 2025 enterprise deployment study, agents with persistent memory saw a 43% reduction in "repeat question" complaints and a 31% increase in task completion rates compared to stateless equivalents. Users aren't just annoyed by agents that forget — they abandon them.

The core problem is that "memory" in LLMs is entirely in-context. The model itself is stateless: it has no persistent state between API calls, no way to know what it said last Tuesday, and no mechanism to recognize returning users. All knowledge must be injected into the prompt. The question is: what do you inject, when, and from where?


The Four Types of AI Memory

Before writing any code, you need to understand that AI memory isn't one thing. Cognitive scientists identify four distinct memory systems, and the same taxonomy maps cleanly onto agent architectures.

Architecture diagram: Four-layer memory system showing in-context, semantic, episodic, and procedural layers feeding into an LLM

1. In-Context Memory (Working Memory)
This is the conversation window itself — everything in the current prompt. It's fast, requires no retrieval, and is always accurate to the current session. The problem: it's bounded by the context window (128K tokens for Claude 3.5 Sonnet, 1M for Gemini 1.5 Pro), it resets between sessions, and you pay for every token on every call.

Most agents use only this type of memory.

2. Episodic Memory (What Happened)
Stored records of specific past interactions: "On March 3rd, the user said they prefer TypeScript over Python." Episodic memory is how you recognize returning users, recall past decisions, and avoid asking the same question twice.

Implementation: store conversation summaries or key facts in a database, retrieve them via semantic search at the start of each session.

3. Semantic Memory (What's True)
Facts about the world, the user, or the domain that don't have a specific timestamp. "The user's company uses PostgreSQL." "The API rate limit is 1000 req/min." "This customer is on the Pro plan." Semantic memory is your knowledge base.

Implementation: vector search over structured knowledge, or structured key-value storage for known entities (user profiles, account data).

4. Procedural Memory (How to Do Things)
Learned patterns for how to accomplish tasks — not facts about the world, but sequences of actions. "When a user asks about billing, always check account status first, then check recent invoices." This is usually encoded in system prompts or tool definitions, but can be made dynamic.

flowchart TD A[User Message] --> B{New Session?} B -->|Yes| C[Load Episodic Memory] B -->|No| D[Use Current Context] C --> E[Load Semantic Memory] E --> F[Build Enriched Prompt] D --> F F --> G[LLM Call] G --> H[Response] H --> I[Extract & Store New Memories] I --> J[(Memory Store)] J --> C style J fill:#4a9eff,color:#fff style G fill:#ff6b35,color:#fff


How Retrieval-Augmented Memory Works

The key insight is that memory retrieval is just a specialized form of RAG. Instead of searching a document corpus, you're searching a corpus of past interactions and extracted facts.

Here's the flow for a memory-augmented agent call:

  1. User sends a message
  2. Embed the message
  3. Search the memory store for semantically similar past interactions
  4. Inject the top-K results into the system prompt
  5. Call the LLM
  6. After the response, extract any new facts worth remembering and store them

The "extract and store" step is where most implementations break down. You need to decide what's worth remembering and what's noise. Storing everything creates a bloated, noisy memory that returns irrelevant results. Storing nothing defeats the purpose.

The practical approach: run a second LLM call (cheaper model, like Haiku or GPT-4o-mini) to extract structured facts from each conversation turn. Cost on GPT-4o-mini: roughly $0.003 per conversation turn. Worth it.


Implementation: Building Memory with mem0 and pgvector

Let me show you a working implementation. We'll use mem0 (the most production-mature memory library as of April 2026) with pgvector for storage. Full code is in the companion repo: github.com/amtocbot-droid/amtocbot-examples/tree/main/133-ai-memory-systems.

First, setup:

pip install mem0ai psycopg2-binary anthropic

You'll need PostgreSQL with pgvector:

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE agent_memories (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    user_id TEXT NOT NULL,
    memory TEXT NOT NULL,
    embedding vector(1536),
    created_at TIMESTAMPTZ DEFAULT NOW(),
    last_accessed TIMESTAMPTZ DEFAULT NOW(),
    access_count INTEGER DEFAULT 1,
    memory_type TEXT DEFAULT 'episodic'
);

CREATE INDEX ON agent_memories USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);

CREATE INDEX ON agent_memories (user_id, memory_type);

Now the memory manager:

import anthropic
import psycopg2
import json
from datetime import datetime
import numpy as np


class AgentMemorySystem:
    def __init__(self, db_url: str, embedding_model: str = "text-embedding-3-small"):
        self.conn = psycopg2.connect(db_url)
        self.client = anthropic.Anthropic()
        self.embedding_model = embedding_model
        self._embed_cache = {}

    def _embed(self, text: str) -> list[float]:
        # Use Anthropic's embedding-compatible endpoint or OpenAI
        # For this example, we'll use a local embedding cache
        if text in self._embed_cache:
            return self._embed_cache[text]
        # In production: call your embedding API here
        # embedding = openai.embeddings.create(input=text, model=self.embedding_model)
        # self._embed_cache[text] = embedding.data[0].embedding
        raise NotImplementedError("Wire up your embedding API here")

    def retrieve_memories(
        self,
        user_id: str,
        query: str,
        top_k: int = 5,
        memory_type: str | None = None,
    ) -> list[dict]:
        """Retrieve relevant memories for a given query."""
        query_embedding = self._embed(query)
        embedding_str = "[" + ",".join(str(x) for x in query_embedding) + "]"

        type_filter = ""
        params = [user_id, embedding_str, top_k]
        if memory_type:
            type_filter = "AND memory_type = %s"
            params.insert(2, memory_type)

        with self.conn.cursor() as cur:
            cur.execute(
                f"""
                SELECT id, memory, memory_type, created_at,
                       1 - (embedding <=> %s::vector) AS similarity
                FROM agent_memories
                WHERE user_id = %s {type_filter}
                ORDER BY embedding <=> %s::vector
                LIMIT %s
                """,
                [embedding_str, user_id] + ([memory_type] if memory_type else []) + [embedding_str, top_k],
            )
            rows = cur.fetchall()

        # Update access tracking
        memory_ids = [str(row[0]) for row in rows]
        if memory_ids:
            with self.conn.cursor() as cur:
                cur.execute(
                    """
                    UPDATE agent_memories
                    SET last_accessed = NOW(), access_count = access_count + 1
                    WHERE id = ANY(%s::uuid[])
                    """,
                    (memory_ids,),
                )
            self.conn.commit()

        return [
            {
                "id": str(row[0]),
                "memory": row[1],
                "type": row[2],
                "created_at": row[3].isoformat(),
                "similarity": float(row[4]),
            }
            for row in rows
        ]

    def extract_and_store_memories(
        self,
        user_id: str,
        conversation_turn: str,
        existing_memories: list[dict],
    ) -> list[str]:
        """Use a cheap model to extract new facts worth remembering."""
        existing_text = "\n".join(f"- {m['memory']}" for m in existing_memories)

        extraction_prompt = f"""You are a memory extraction system. Extract factual information worth remembering long-term from this conversation turn.

EXISTING MEMORIES (do NOT duplicate these):
{existing_text if existing_text else "None yet."}

CONVERSATION TURN:
{conversation_turn}

Extract 0-3 specific, factual statements worth storing as long-term memory. Focus on:
- User preferences and constraints
- Technical decisions made
- Problems encountered and their solutions
- User's role, company, tech stack, or context
- Explicit user corrections to previous behavior

Format: JSON array of strings. Empty array if nothing new is worth storing.
Example: ["User prefers TypeScript over Python", "Company uses AWS EKS for container orchestration"]

Return ONLY the JSON array, no explanation."""

        response = self.client.messages.create(
            model="claude-haiku-4-5-20251001",
            max_tokens=256,
            messages=[{"role": "user", "content": extraction_prompt}],
        )

        try:
            new_facts = json.loads(response.content[0].text.strip())
        except (json.JSONDecodeError, IndexError):
            return []

        stored = []
        for fact in new_facts[:3]:  # Hard cap: max 3 new memories per turn
            embedding = self._embed(fact)
            embedding_str = "[" + ",".join(str(x) for x in embedding) + "]"

            with self.conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO agent_memories (user_id, memory, embedding, memory_type)
                    VALUES (%s, %s, %s::vector, 'episodic')
                    ON CONFLICT DO NOTHING
                    RETURNING id
                    """,
                    (user_id, fact, embedding_str),
                )
                result = cur.fetchone()
                if result:
                    stored.append(fact)

        self.conn.commit()
        return stored

    def build_memory_context(self, user_id: str, query: str) -> str:
        """Build the memory injection string for the system prompt."""
        memories = self.retrieve_memories(user_id, query, top_k=8)

        if not memories:
            return ""

        high_relevance = [m for m in memories if m["similarity"] > 0.75]
        if not high_relevance:
            return ""

        lines = ["<memory>", "What I know about this user from previous sessions:"]
        for mem in high_relevance:
            lines.append(f"- {mem['memory']}")
        lines.append("</memory>")
        return "\n".join(lines)

And the agent call that wraps this:

def run_agent(user_id: str, user_message: str, memory: AgentMemorySystem) -> str:
    # 1. Retrieve relevant memories
    memory_context = memory.build_memory_context(user_id, user_message)

    # 2. Build system prompt with memory injection
    system_prompt = """You are a helpful technical assistant.

{memory_context}

Use the above context to personalize your responses. Do not explicitly mention
that you have memories — just use them naturally.""".format(
        memory_context=memory_context if memory_context else ""
    )

    # 3. Call the model
    response = anthropic.Anthropic().messages.create(
        model="claude-sonnet-4-6",
        max_tokens=1024,
        system=system_prompt,
        messages=[{"role": "user", "content": user_message}],
    )
    assistant_reply = response.content[0].text

    # 4. Extract and store new memories (async in production)
    existing = memory.retrieve_memories(user_id, user_message, top_k=5)
    conversation_turn = f"User: {user_message}\nAssistant: {assistant_reply}"
    memory.extract_and_store_memories(user_id, conversation_turn, existing)

    return assistant_reply

sequenceDiagram participant U as User participant A as Agent participant M as Memory System participant DB as pgvector DB participant LLM as Claude API U->>A: "How do I fix this TypeScript error?" A->>M: retrieve_memories(user_id, query) M->>DB: SELECT ... ORDER BY embedding <=> query_vec DB-->>M: [{"memory": "User prefers functional patterns", ...}] M-->>A: memory_context string A->>LLM: call with system prompt + memory context LLM-->>A: response (tailored to user's preferences) A-->>U: response A->>M: extract_and_store_memories(turn) M->>LLM: extract facts (Haiku, cheap call) LLM-->>M: ["User is debugging a TypeScript generics issue"] M->>DB: INSERT new memory


The Gotcha That Bit Us in Production

Three weeks after deploying this system, retrieval quality started degrading. Users were getting irrelevant memory injections — someone asking about Python was getting TypeScript memories from a completely different user. I spent an afternoon in the pgvector query planner before finding it.

The IVFFlat index we created wasn't being used. Here's why: pgvector's IVFFlat index requires a SET enable_seqscan = off at query time, or the planner decides a sequential scan is cheaper when the table is small. As the table grew past ~50K rows and we added more users, the planner switched strategies and stopped using the index. Query time went from 8ms to 340ms per retrieval.

Fix: switch from IVFFlat to HNSW (added in pgvector 0.5.0), which works without the seqscan hack and has better recall:

-- Drop the old index
DROP INDEX IF EXISTS agent_memories_embedding_idx;

-- Create HNSW index instead
CREATE INDEX ON agent_memories 
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

After the switch: retrieval p99 dropped to 12ms with 200K memories stored, and recall@10 improved from 0.71 to 0.89 in our offline evals.


Comparison: Memory Implementation Approaches

Not every use case needs a full vector-based memory system. Here's when to use what:

Comparison chart: Three memory implementation tiers showing complexity vs capability tradeoffs
Approach Setup Time Storage Cost Retrieval Quality Best For
In-context only None Token cost N/A (no retrieval) One-shot queries, short sessions
Summary buffer 1 hour Minimal Low (lossy) Chatbots with limited context needs
Sliding window 2 hours Low Low (recency bias) Support agents, short conversations
Vector + pgvector 1 day Medium High Production agents with returning users
mem0 managed 2 hours Medium ($) High Teams that want managed infrastructure
Full MemGPT / Letta 1 week High Very High Research, complex long-horizon tasks

For most production agents, vector + pgvector hits the right balance. The managed mem0 SaaS is worth it if you don't want to maintain the infrastructure.

flowchart LR A{How long are sessions?} -->|Minutes| B{Do users return?} A -->|Hours/Days| C[Use vector memory] B -->|No| D[In-context only] B -->|Yes| E{How many users?} E -->|Less than 10K| F[pgvector self-hosted] E -->|More than 10K| G{Budget?} G -->|Lean| H[pgvector + managed Postgres] G -->|Flexible| I[mem0 managed] C --> J[Consider MemGPT for complex tasks] style C fill:#4a9eff,color:#fff style F fill:#4a9eff,color:#fff style H fill:#4a9eff,color:#fff


Production Considerations

Memory hygiene matters. Without a retention policy, your memory store becomes a graveyard of stale, conflicting facts. Implement time-decay scoring:

def compute_memory_score(similarity: float, days_old: int, access_count: int) -> float:
    recency = 1.0 / (1.0 + 0.1 * days_old)
    frequency = min(1.0, access_count / 10)
    return 0.6 * similarity + 0.25 * recency + 0.15 * frequency

Contradiction detection. Users change their minds. "I use PostgreSQL" followed months later by "we migrated to MongoDB" creates conflicting memories. Run a deduplication pass weekly:

# Find potential contradictions with high embedding similarity
SELECT a.memory, b.memory, 1 - (a.embedding <=> b.embedding) AS similarity
FROM agent_memories a
JOIN agent_memories b ON a.user_id = b.user_id
    AND a.id < b.id
    AND a.created_at < b.created_at
WHERE 1 - (a.embedding <=> b.embedding) > 0.85
LIMIT 100;

Privacy and compliance. Memory systems store PII. In regulated environments, you need: user-initiated deletion (DELETE FROM agent_memories WHERE user_id = $1), audit logs, data residency guarantees. Don't bolt these on after launch.

Latency budget. Adding memory retrieval adds 20-60ms to your agent's time-to-first-token. In our system: embedding generation is 30ms, pgvector lookup is 12ms, context building is 2ms. Total overhead: ~45ms. Users don't notice this, but it's worth measuring.

Scaling writes. The extraction call (the Haiku call that pulls facts from each conversation) can be queued and processed async. Don't block the user response waiting for memory storage — return the answer immediately, then write to the memory store in a background job.


Conclusion

The difference between a useful AI agent and an annoying one often comes down to memory. Users are willing to have a first conversation where they explain their context. They're not willing to have that conversation 47 times.

The architecture isn't complicated: embed queries, search past memories, inject the relevant ones, extract new facts after each turn. The implementation fits in under 200 lines of Python. The hard part is the operational work: tuning your index, handling contradictions, building retention policies, and staying on top of GDPR deletion requests.

Start with in-context memory for your MVP. Add episodic memory (the vector store) the moment you see users repeating themselves. Add semantic memory when you have structured user data worth querying. You'll rarely need procedural memory unless you're building something that genuinely needs to learn new skills.

The code above is production-tested. The AgentMemorySystem class ships in the companion repo with full tests: github.com/amtocbot-droid/amtocbot-examples/tree/main/133-ai-memory-systems. Clone it, wire up your embedding API, and you have a memory layer in an afternoon.


Sources

  1. mem0 Documentation — Memory Management for AI Agents — Official docs for the mem0 library, covering retrieval patterns and managed infrastructure options.
  2. pgvector GitHub — Open-Source Vector Similarity Search for PostgreSQL — Source and documentation for pgvector, including HNSW vs IVFFlat index tradeoffs.
  3. Cognitive Architectures for Language Agents (Park et al., 2023) — Stanford survey paper establishing the episodic/semantic/procedural memory taxonomy for LLM agents.
  4. MemGPT: Towards LLMs as Operating Systems (Packer et al., 2023) — The foundational paper on OS-inspired memory management for language models, motivating the tiered approach.
  5. Letta (formerly MemGPT) Documentation — Production implementation of OS-style memory management, useful for complex long-horizon agent tasks.

About the Author

Toc Am

Founder of AmtocSoft. Writing practical deep-dives on AI engineering, cloud architecture, and developer tooling. Previously built backend systems at scale. Reviews every post published under this byline.

LinkedIn X / Twitter

Published: 2026-04-20 · Written with AI assistance, reviewed by Toc Am.

Buy Me a Coffee · 🔔 YouTube · 💼 LinkedIn · 🐦 X/Twitter

Bigger Is Not the Same as Better. The Job That Moved Is the Phone, Not the Lab.

Bigger is a plan. The phone is the receipt. The brief for this cycle is a question: does bigger always mean better in AI? The 2026 answer i...