Memory Architecture with Lakebase and MLflow · Topic 2 of 3
Retrieving Memories from Lakebase
Databricks Context Engineer: choose AI Search or SQL for memory lookups, weigh static against dynamic retrieval, and diagnose over and under retrieval.
- 5 min
- Meaning, paraphrase or similarity: AI Search. Exact keys, filters, counts, "latest": structured SQL.
- Semantic search returns top-k, so it cannot guarantee "all" or an exact count.
- Static retrieval is predictable and fast but costs tokens every turn and can go stale.
What it is
Storing memories is half the job; the other half is bringing the right ones back. Memory retrieval decides three things: how a memory is found (by meaning or by exact match), when it is loaded (once up front or on demand) and how much enters the window.
How it works
AI Search or structured query
Memories persisted in Lakebase sit in Postgres tables with columns such as user ID, memory type, timestamps and the memory text. Some questions are about meaning; others are about exact values.
| Query type | Mechanism | Why |
|---|---|---|
| "What has this customer said about delivery problems before?" | Databricks AI Search (vector search over memory content) | Wording varies, so match by meaning |
| "Find past conversations similar to this complaint" | AI Search | Similarity is the goal |
| "What is the user's current plan tier?" | Structured SQL query on Lakebase | One exact value by key |
| "List all approvals recorded after 1 October" | Structured SQL query | Exact filter on a date column |
| "How many open escalations does this account have?" | Structured SQL query | A count needs exact rows, not top-k |
| "Past notes about refunds, for this user only" | AI Search with a filter on user ID | Meaning, scoped by an exact column |
Semantic search returns the top k most similar items, so it cannot reliably answer "all", "how many" or "the latest". A key lookup cannot match a paraphrase.
SELECT plan_tier, updated_at
FROM agent_memory.user_profile
WHERE user_id = 'u-2291';Static versus dynamic retrieval
Static retrieval loads a fixed set of memories into the context before the model call, usually at the start of the session. Dynamic retrieval gives the agent a memory tool that it calls mid-run, only when a step needs memory.
| Static (load up front) | Dynamic (agent calls a tool) | |
|---|---|---|
| Latency | One lookup, no extra tool turns | A tool round trip each time memory is needed |
| Predictability | Same context every time, easy to test | Depends on the agent deciding to call the tool |
| Token cost | Paid on every turn, used or not | Paid only when needed |
| Freshness | Can go stale if memory changes mid-session | Reads current state at the moment of need |
| Fit | Small, always relevant memory such as a profile | Large memory where only some items matter per step |
Worked numbers: static loading of 20 memories at about 150 tokens each puts 3,000 tokens in every turn, or 30,000 tokens over a 10-turn session. If only two turns need memory and each dynamic call returns 5 memories, the cost is about 1,500 tokens plus two tool round trips.
A common hybrid loads a small static core (identity, plan, key preferences) and lets the agent retrieve the rest dynamically.
Over-retrieval and under-retrieval
Over-retrieval fills the window with memories the step does not need. This is context pollution: irrelevant items cause confusion, an outdated memory next to a new one causes a clash, and the extra tokens raise cost and latency.
Under-retrieval leaves out history the step does need. The agent asks again for information it was given, ignores a stated constraint or contradicts an earlier decision.
| Signal in the trace | Diagnosis | Remediation |
|---|---|---|
| 25 memories retrieved, most unrelated to the question | Over-retrieval | Lower the top-k, add filters (user, type, recency), set a relevance threshold |
| An old and a corrected preference both retrieved | Over-retrieval of superseded items | Mark or delete superseded memories on update |
| Memory search returned nothing relevant, answer misses known history | Under-retrieval | Rewrite the query, use hybrid keyword plus vector search, widen a filter that is too narrow |
| The memory tool was never called on a step that needed it | Under-retrieval from dynamic loading | Make the tool description state when to use it, or load that item statically |
In MLflow 3 evaluation, retrieval relevance scorers flag polluting items and correctness against expected facts reveals missing history.
How the exam asks about this
- "Count how many orders a user returned last quarter from stored memories." Use a structured SQL query on Lakebase, not semantic search.
- "A shopping agent loads 40 past memories every turn and starts recommending products the user rejected long ago." Over-retrieval causing pollution and clash; filter, cap and supersede.
- "Memory is large but each step needs only a few items, and memory changes during the session." Prefer dynamic retrieval through a memory tool.
What to remember for the exam
- Meaning, paraphrase or similarity: AI Search. Exact keys, filters, counts, "latest": structured SQL.
- Semantic search returns top-k, so it cannot guarantee "all" or an exact count.
- Static retrieval is predictable and fast but costs tokens every turn and can go stale.
- Dynamic retrieval is targeted and fresh but adds tool calls and depends on the agent calling the tool.
- Over-retrieval pollutes context; under-retrieval loses needed history.
- Fix over-retrieval with filters, lower top-k and superseding old items; fix under-retrieval with better queries and wider recall.
Raising top-k is not a general cure for missing history. If the query or filter is wrong, more results add pollution without finding the needed memory.
Common mistakes
Using vector search for exact facts. "The user's account ID" or "the number of open tickets" needs a key lookup or a count, not the five most similar memories.
Loading everything statically to be safe. Every turn pays for all of it, and stale or rejected items compete with the current request.
Blaming the model for forgetting. When the agent ignores known history, check the trace first: the retrieval may have returned nothing or never run.
Check yourself
Which mechanism finds every memory created after a given date for one user?
A structured SQL query on the Lakebase table, because it needs an exact date filter and all matching rows, not top-k similarity.
What is the main cost of static memory retrieval?
The loaded memories are paid for in tokens on every turn whether used or not, and they can go stale if memory changes during the session.
An agent repeats a suggestion the user rejected, and the trace shows both the old and the new preference in context. What is this?
Over-retrieval causing context pollution and clash; supersede the old memory and filter retrieval so only current items are returned.
Topic 2 of 3 in Memory Architecture with Lakebase and MLflow
Test yourself on Memory Architecture with Lakebase and MLflow
The full Databricks Context Engineer Associate question bank has 600 questions, with timed practice exams and an explanation for every answer.
This study guide is independent and is not affiliated with, endorsed by, or sponsored by Databricks. Vendor names, certification names and exam codes are trademarks of their respective owners and are used for identification only.