Course Content
Context Management and Memory
3 sections · 6 lessons
Retrieval, Caching, and Embeddings
A product team's internal assistant has a memory store with 4,000 entries. Someone asks it: "what did we decide about the pricing page?"
It retrieves the top 5 by cosine similarity and answers with a summary of a design discussion about button placement. It never mentions the actual decision — three tiers, annual toggle on by default, signed off on 2 April — even though that memory is sitting in the store, correctly written, correctly embedded.
The retrieval log explains it in one glance:
rank score memory 1 0.88 "In the pricing page review we discussed the layout of the tier cards..." 2 0.86 "...continuing the pricing page layout discussion, the tier cards should..." 3 0.85 "...on the pricing page, card ordering and the tier card headings were..." 4 0.71 "Decided 2026-04-02: pricing page ships with three tiers, annual toggle default on." 5 0.68 "Pricing page A/B test concluded; variant B won on conversion by 4.1%."Ranks 1, 2 and 3 are three overlapping fragments of the same forty-minute meeting. They are each highly similar to the query and near-identical to each other. Taking the top 3 fills the entire memory budget with one conversation and pushes the answer to rank 4, outside the cut.
Nothing failed. The embeddings are good, the index is correct, the scores are right. Plain top-k similarity simply optimises for the wrong thing: it maximises relevance per result and completely ignores redundancy across results.
Two failures, not one
Retrieval systems fail in two independent ways, and they need different fixes.
| Quality failure | Performance failure | |
|---|---|---|
| Symptom | Wrong or redundant results; the right memory exists but is not returned | Slow responses; high bills; rate-limit errors |
| Where it shows up | User complaints about accuracy | p95 latency graph; monthly invoice |
| Fixed by | Diversity (MMR), filtering, score floors, evaluation | Caching, batching, dimension reduction |
| Danger if ignored | Users stop trusting the assistant | Costs grow linearly with usage until someone notices |
Similarity search and its blind spot
The baseline is: embed the query, find the k nearest stored vectors, return them. It is the right default and it is one line of code.
results = store.query(query_texts=[user_message], n_results=5, where={"user_id": user_id})Its blind spot is that the k results are chosen independently. If three memories are near-duplicates of each other, all three will be returned, because each one individually scores well against the query. Nothing in the algorithm notices that you have now spent 60% of your slots on one idea.
Top-k similarity answers "which memories are most like this query?" It never asks "given what I have already selected, what does this add?" That second question is the one that decides whether the answer is complete.
Read the scores, do not just take the top k
Most vector stores return distances, not similarities, and the conversion depends on the metric configured on the index. Get this wrong and every threshold you set is nonsense.
1res = store.query(query_texts=[q], n_results=10, where={"user_id": uid},2 include=["documents", "metadatas", "distances"])34scored = []5for doc, meta, dist in zip(res["documents"][0], res["metadatas"][0],6 res["distances"][0]):7 similarity = 1.0 - dist # valid for COSINE distance only8 scored.append((similarity, doc, meta))For cosine distance, similarity is 1−d. For squared Euclidean distance on normalised vectors it is 1−d/2. For inner-product indices the raw score is the similarity. Log a handful of known-good and known-bad pairs on day one and confirm that good pairs score high — if your "similar" pairs come out near zero, you have the conversion backwards.
Calibrating a score floor
A floor discards results below a threshold rather than always returning k things. To set it, label a few dozen query–memory pairs by hand and look at the two distributions.
| Floor | Avg memories injected | Of which relevant | Precision | Relevant items missed |
|---|---|---|---|---|
| none (always top-5) | 5.0 | 1.55 | 31% | 0.00 |
| 0.30 | 4.2 | 1.55 | 37% | 0.00 |
| 0.45 | 2.1 | 1.48 | 70% | 0.07 |
| 0.60 | 1.1 | 1.02 | 93% | 0.53 |
Moving the floor from none to 0.45 more than doubles precision (31% to 70%) while losing 0.07 relevant items per query — a very good trade. Pushing on to 0.60 gains another 23 points of precision but now misses half a relevant memory per query, which is how an assistant starts saying "I don't have that information" about things it definitely knows.
The right floor depends on the cost asymmetry. For a medical or financial assistant, injecting an irrelevant memory is worse than missing one, so go high. For a casual assistant, missing context is the more visible failure, so go low.
Maximal Marginal Relevance
MMR fixes the redundancy blind spot by selecting results one at a time, each time penalising similarity to what has already been chosen.
where R is the candidate pool, S is what you have selected so far, q is the query, and λ trades relevance against diversity. Work the pricing-page example with λ=0.5. The pairwise similarities among the candidates are:
D1 D2 D3 D4 D5 sim to query D1 - 0.94 0.92 0.31 0.25 0.88 D2 0.94 - 0.93 0.29 0.28 0.86 D3 0.92 0.93 - 0.27 0.26 0.85 D4 0.31 0.29 0.27 - 0.40 0.71 D5 0.25 0.28 0.26 0.40 - 0.68Round 1. Nothing is selected yet, so the highest query similarity wins: D1 at 0.88. Now S={D1}.
Round 2. Score each remaining candidate as 0.5⋅sim(d,q)−0.5⋅sim(d,D1):
- D2: 0.5(0.86)−0.5(0.94)=0.430−0.470=−0.040
- D3: 0.5(0.85)−0.5(0.92)=0.425−0.460=−0.035
- D4: 0.5(0.71)−0.5(0.31)=0.355−0.155=+0.200
- D5: 0.5(0.68)−0.5(0.25)=0.340−0.125=+0.215
D5 wins. The two near-duplicates went negative — their redundancy penalty exceeded their relevance contribution. Now S={D1,D5}.
Round 3. The penalty is now the maximum similarity to either selected item:
- D2: 0.430−0.5max(0.94,0.28)=0.430−0.470=−0.040
- D3: 0.425−0.5max(0.92,0.26)=0.425−0.460=−0.035
- D4: 0.355−0.5max(0.31,0.40)=0.355−0.200=+0.155
D4 wins. Final selection: D1, D5, D4 — the meeting discussion, the A/B result, and the decision. Plain similarity would have returned D1, D2, D3: the same meeting three times, and no decision.
| λ | Behaviour | Use for |
|---|---|---|
| 1.0 | Pure similarity — MMR disabled | Precise factual lookup where duplicates are impossible |
| 0.7 | Mild diversity nudge | Sensible default for memory retrieval |
| 0.5 | Balanced | Summarising a topic; "what do we know about X" |
| 0.3 | Diversity dominates | Exploratory browsing; risks surfacing weakly relevant items |
The practical shape is to over-fetch then re-rank: pull 20 candidates by similarity, run MMR to pick 5. MMR over the whole store would be quadratic and far too slow; over 20 candidates it is 20 × 5 comparisons and takes microseconds.
Filtering: before, not after
Metadata filters restrict the search space by user, kind, date or source. The critical detail is where the filter is applied.
| Pre-filter (in the query) | Post-filter (in your code) | |
|---|---|---|
| What happens | Index searches only matching rows, returns 5 | Index searches everything, returns 5, you discard non-matching ones |
| Results after filtering | Always 5 | Anywhere from 0 to 5 |
| Correctness | Correct | Silently drops results; on a busy store often returns nothing |
| Security | Other users' data never leaves the database | Other users' data is loaded into your process, one bug away from a leak |
The arithmetic is unforgiving. With 4,000 memories across 200 users, one user owns about 20 of them — 0.5% of the store. A post-filtered top-5 search over the whole store returns, on average, 5×0.005=0.025 of that user's memories. You will get zero results roughly 97 times out of 100, and the three times you get something, it belonged to someone else until your code threw it away.
1# Correct: the database never considers other users' rows.2res = store.query(3 query_texts=[q],4 n_results=20,5 where={"$and": [6 {"user_id": {"$eq": user_id}},7 {"kind": {"$in": ["fact", "decision", "constraint"]}},8 {"created_at": {"$gte": six_months_ago}},9 ]},10)Where retrieval fits in the whole call
Retrieval into a prompt is the same machinery as retrieval-augmented generation over documents; the only difference is that the corpus is the user's own history rather than a knowledge base. The assembled request looks like this:
system prompt 3,000 tok cacheable, byte-stableretrieved memories (5 x 40) 200 tok changes every turnrecent verbatim turns 4,800 tok changes every turncurrent user message 120 tok --------- 8,120 tok inputNote the ordering. Everything stable sits at the front, everything volatile at the back. That is not cosmetic — it is what makes prompt caching work, and prompt caching is the largest single cost lever in the whole system.
Caching, ranked by what it actually saves
There are four things worth caching and they are wildly unequal. Take a service handling 1,000,000 requests per month, each with the token profile above, at 3 dollars per million input tokens and 15 per million output.
Baseline cost per request: input 8,120×3/106=0.02436 dollars, output 320×15/106=0.0048 dollars, total about 2.92 cents, so roughly 29,160 dollars per month.
| Cache | What it stores | Realistic hit rate | Monthly saving |
|---|---|---|---|
| Prompt cache (provider-side) | The 3,000-token stable prefix | 85% | ~6,550 dollars |
| Response cache | Full replies to identical requests | 12% | ~3,499 dollars |
| Retrieval cache | Result sets for a (user, query) pair | 25% | ~0 dollars; ~4 ms per hit |
| Embedding cache | Vectors for repeated query text | 35% | ~0 dollars; ~28 ms per hit |
The prompt-cache figure: the prefix costs 3,000×3/106=0.009 dollars uncached and about a tenth of that when read from cache, so 0.0081 dollars saved per hit × 1,000,000 × 0.85 = 6,885 dollars. Each of the 150,000 misses writes the prefix back to the cache at 1.25 times the input price, an extra 0.25×0.009=0.00225 dollars, or about 340 dollars in total, so the net saving is about 6,550 dollars. The response-cache figure: 0.02916 × 1,000,000 × 0.12 = 3,499 dollars, less the cost of running the cache itself.
Ordering the stable part of your prompt first and enabling prompt caching is worth more than every application-level cache combined, and it takes one line of configuration.
Retrieval and embedding caches save almost no money, because embedding is priced in cents per million tokens and vector search costs nothing per query. Build them anyway — but build them for latency and dependency isolation, not for the invoice. When your embedding provider has a bad ten minutes, a warm cache is the difference between degraded and down.
An in-process cache
1from functools import lru_cache2import hashlib34@lru_cache(maxsize=10_000)5def embed_cached(text: str) -> tuple:6 return tuple(embedding_model.encode(text)) # tuple: hashable, immutableFine for a single process. Useless the moment you run three replicas, because each holds its own copy and the effective hit rate falls to roughly a third of what a shared cache would give.
Redis, and the invalidation problem nobody plans for
The hard part of caching retrieval results is not storing them. It is knowing when they are wrong. The instant a user adds a memory, every cached result set for that user is potentially stale — and enumerating and deleting those keys is slow and error-prone.
The clean solution is a per-user version counter embedded in the cache key. Writing a memory bumps the counter, which makes every old key unreachable in one atomic operation. Stale entries are never read again and expire on their own TTL.
1import json, hashlib, redis23r = redis.Redis(decode_responses=True)45def _version(user_id: str) -> int:6 return int(r.get(f"memver:{user_id}") or 0)78def invalidate(user_id: str) -> None:9 r.incr(f"memver:{user_id}") # O(1); no key scanning1011def cached_recall(user_id: str, query: str, k: int = 5, ttl: int = 900):12 h = hashlib.sha256(f"{query}|{k}".encode()).hexdigest()[:16]13 key = f"ret:{user_id}:v{_version(user_id)}:{h}"1415 hit = r.get(key)16 if hit is not None:17 return json.loads(hit)1819 results = recall(user_id, query, k=k)20 r.setex(key, ttl, json.dumps(results))21 return results| Cache | TTL | Why |
|---|---|---|
| Embeddings | Days, or until the embedding model changes | The same text always embeds to the same vector |
| Retrieval results | 5–15 minutes, plus version-key invalidation | New memories should show up quickly |
| Full responses | Minutes to an hour, non-personalised only | Answers age; personalised answers must never be shared |
The semantic-cache trap
A tempting idea: cache responses by embedding similarity, so "how do I reset my password" hits the entry stored for "password reset steps". It works, and then it silently breaks.
"What is my refund amount?" embeds at 0.96 similarity to"What is my refund status?"Different questions. Different answers. One cached response.Two rules make semantic caching safe. First, apply it only to non-personalised, generic queries — never to anything whose answer depends on who is asking. Second, set the threshold high, around 0.97, and accept a low hit rate; at 0.90 you will serve wrong answers to a measurable fraction of users, and because they are plausible answers nobody will report them.
Measuring whether retrieval is any good
You cannot tune what you do not measure, and "it seems better" is not measurement. Build a small labelled set — 50 queries with the memory IDs that should be retrieved for each — and compute three numbers.
Take two queries from such a set:
Q1 "what are the user's dietary constraints?" relevant = {m17, m43} retrieved top-5 = [m43, m88, m17, m12, m91] precision@5 = 2/5 = 0.400 recall@5 = 2/2 = 1.000 first relevant at rank 1 -> RR = 1/1 = 1.000Q2 "which laptop did they buy?" relevant = {m5} retrieved top-5 = [m22, m31, m5, m9, m14] precision@5 = 1/5 = 0.200 recall@5 = 1/1 = 1.000 first relevant at rank 3 -> RR = 1/3 = 0.333Averaged: precision@5 = 0.300, recall@5 = 1.000, MRR = (1.000+0.333)/2=0.667.
| Metric | Question it answers | Act on it when |
|---|---|---|
| Recall@k | Did the right memory make the cut at all? | Low → raise k, improve extraction, check the embedding model |
| Precision@k | How much of what I inject is noise? | Low → raise the floor, add MMR, tighten filters |
| MRR | How near the top is the first useful result? | Low with good recall → a re-ranking problem, not a retrieval problem |
Recall is the one to protect. If the right memory is not in the candidate pool, no amount of re-ranking will conjure it. Precision can always be recovered later with a floor or a re-ranker.
Edge cases that decide how the system feels
Nothing found
The wrong response is to inject a block saying "No relevant memories found." Models read that as content and will happily tell the user "I searched my memory and found nothing", which sounds broken. The right response is to inject nothing at all and let the assistant behave like an assistant meeting the topic for the first time.
hits = recall(user_id, query, k=5, floor=0.45)memory_block = format_memories(hits) if hits else "" # empty, not a messageLow confidence
Results between the floor and a comfortable threshold — say 0.45 to 0.55 — are worth including but not worth asserting. Label them so the model hedges rather than states.
1def format_memories(hits):2 strong = [h for h in hits if h["score"] >= 0.55]3 weak = [h for h in hits if h["score"] < 0.55]4 parts = []5 if strong:6 parts.append("Known about this user:\n" +7 "\n".join(f"- {h['text']}" for h in strong))8 if weak:9 parts.append("Possibly relevant, confirm before relying on it:\n" +10 "\n".join(f"- {h['text']}" for h in weak))11 return "\n\n".join(parts)An assistant that says "I think you mentioned you're based in Berlin — is that still right?" is far more trustworthy than one that either asserts it flatly or says nothing.
Getting this right in a real service
Over-fetch, then re-rank. The single highest-value change to a naive retrieval pipeline is to pull 20 candidates instead of 5 and select the final 5 with MMR and a score floor. It costs nothing measurable — the index does the same work — and it fixes the redundancy failure that produced the pricing-page answer.
Put the filter in the query. Not in a list comprehension afterwards. Write an integration test that creates two users, writes a distinctive memory for each, and asserts that a query as user A never returns user B's row. That test is thirty lines and it is the difference between a feature and an incident report.
Order the prompt for caching. Stable prefix first, volatile retrieved content last. Then check cache_read_input_tokens on a real response. If it is zero, something in your prefix is changing between calls — a timestamp, a request ID, a tool list built from an unordered dictionary — and you are paying full price on the biggest line item in the budget.
Build the labelled set before you start tuning. Fifty queries with known-correct answers takes an afternoon to assemble and turns every subsequent decision — floor, k, λ, embedding model — into a measurement instead of an argument. Teams that skip this step spend months adjusting thresholds by feel, and the numbers above are exactly the kind of thing you cannot see by feel: a floor change that doubles precision while costing 0.07 relevant items is invisible without a scoreboard.