📣
TiDB Cloud Premium is now in public preview. Unlimited growth, instant elasticity, advanced security for enterprise workloads. Try it out →

Store AI Agent Memory in TiDB



An AI agent that remembers users across sessions needs its memory store to do four things: find memories by meaning, narrow them by user and type, replace an outdated memory without leaving a half-written state, and forget memories when they expire or a user asks. In TiDB, one table and plain SQL cover all four, because vector search, secondary indexes, transactions, and row expiry live in the same database.

This document walks through that table step by step. For a complete chat application built on the same idea with the pytidb Python SDK, see AI Agent Memory Example.

Prerequisites

  • A TiDB Cloud Starter instance. You can create a free TiDB Cloud Starter instance on TiDB Cloud.
  • A MySQL client connected to the instance. In the TiDB Cloud console, navigate to the overview page of your TiDB Cloud Starter instance and click Connect for the connection parameters.

The examples use 3-dimensional vectors with manually assigned values to make the SQL examples and query results easier to understand. These vectors are for demonstration purposes only and do not represent embeddings generated by an embedding model. In your application, set the vector dimension to the output size of your embedding model, for example VECTOR(1536), and generate each embedding with that model.

Step 1. Create the memory table

Select a database first. This example uses the test database that every TiDB Cloud Starter instance has:

USE test; CREATE TABLE agent_memory ( id BIGINT PRIMARY KEY AUTO_RANDOM, user_id VARCHAR(64) NOT NULL, kind ENUM('fact', 'preference', 'episode') NOT NULL, content TEXT NOT NULL, embedding VECTOR(3) NOT NULL, is_active BOOLEAN NOT NULL DEFAULT TRUE, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, KEY idx_user (user_id, is_active), VECTOR INDEX idx_embedding ((VEC_COSINE_DISTANCE(embedding))) ) TTL = `created_at` + INTERVAL 90 DAY TTL_ENABLE = 'ON';

Each part of the table serves one memory operation:

  • embedding and the vector index idx_embedding let you find memories by meaning.
  • user_id and kind let you narrow memories to one user and one type. The secondary index idx_user on (user_id, is_active) makes reading one user's active memories fast.
  • is_active lets you retire an outdated memory without deleting it.
  • The TTL clause makes TiDB automatically clean up memories that have been stored for more than 90 days. Expired memories are not necessarily deleted immediately because TTL cleanup runs periodically. For more information, see TTL (Time to Live).

Step 2. Write memories

INSERT INTO agent_memory (user_id, kind, content, embedding) VALUES ('alice', 'fact', 'Alice is a data engineer at a fintech company.', '[0.9, 0.1, 0.2]'), ('alice', 'preference', 'Alice prefers answers in Python.', '[0.2, 0.9, 0.1]'), ('alice', 'episode', 'Alice asked how to tune a slow join on 2026-10-01.', '[0.5, 0.4, 0.7]'), ('bob', 'fact', 'Bob runs a small online bookstore.', '[0.8, 0.2, 0.3]'), ('bob', 'preference', 'Bob prefers short answers.', '[0.1, 0.8, 0.3]');

Step 3. Replace an outdated memory in a transaction

When a user changes a preference, retire the old memory and write the new one in a single transaction. The transaction ensures that readers do not see a partially applied replacement.

To prevent concurrent writes from creating multiple active preferences for the same user, add a uniqueness constraint as described in Enforce at most one active preference per user.

BEGIN; UPDATE agent_memory SET is_active = FALSE WHERE user_id = 'alice' AND kind = 'preference' AND is_active; INSERT INTO agent_memory (user_id, kind, content, embedding) VALUES ('alice', 'preference', 'Alice now prefers answers in Go.', '[0.25, 0.85, 0.15]'); COMMIT;

Check the result:

SELECT content, is_active FROM agent_memory WHERE user_id = 'alice' AND kind = 'preference' ORDER BY is_active DESC;
+----------------------------------+-----------+ | content | is_active | +----------------------------------+-----------+ | Alice now prefers answers in Go. | 1 | | Alice prefers answers in Python. | 0 | +----------------------------------+-----------+

If a statement in the transaction fails, for example because a value does not fit its column, run ROLLBACK instead of COMMIT. A failed statement does not end the transaction: the statements that succeeded are still pending, and COMMIT would save them, for example leaving the user with the old preference retired and no new one. After ROLLBACK, the memory stays as it was.

Enforce at most one active preference per user

The transaction makes each replacement atomic, but concurrent writes can still create multiple active preferences for the same user. For example, if no active preference exists, two transactions can each update zero rows, insert a new preference, and commit successfully.

If your application allows only one active preference per user, add the following generated column and unique index. For an existing table, first deactivate or remove duplicate active preferences for each user.

ALTER TABLE agent_memory ADD COLUMN active_preference_user VARCHAR(64) GENERATED ALWAYS AS ( CASE WHEN kind = 'preference' AND is_active THEN user_id ELSE NULL END ) VIRTUAL; ALTER TABLE agent_memory ADD UNIQUE KEY uk_active_preference (active_preference_user);

The generated column contains user_id only for active preferences. The unique index prevents two such rows from having the same user ID. Other rows have NULL in this column, so the index allows multiple inactive preferences, facts, and episodes per user.

After adding the unique index, deactivate the old preference and insert the new preference in a single transaction, as shown above. The unique index enforces the limit even under concurrent writes. If the replacement fails, roll back the transaction. To retry, start a new transaction and repeat both the UPDATE and the INSERT.

Step 4. Recall memories for one user

Most agent requests need the memories of the current user only. Filter by user first, then rank by distance to the embedding of the current request:

SELECT kind, content, VEC_COSINE_DISTANCE(embedding, '[0.2, 0.9, 0.2]') AS distance FROM agent_memory WHERE user_id = 'alice' AND is_active ORDER BY distance LIMIT 3;
+------------+----------------------------------------------------+-----------------------+ | kind | content | distance | +------------+----------------------------------------------------+-----------------------+ | preference | Alice now prefers answers in Go. | 0.0032403313698232683 | | episode | Alice asked how to tune a slow join on 2026-10-01. | 0.3295984359391211 | | fact | Alice is a data engineer at a fintech company. | 0.645662235088144 | +------------+----------------------------------------------------+-----------------------+

This query does not use the vector index, because the WHERE filter is applied before the nearest-neighbor search. It can use the secondary index idx_user to read only this user's active memories and computes the exact distance for each one. The result is exact, and the cost primarily depends on the number of matching memories when the secondary index is used, rather than requiring a full-table scan. This is usually the right pattern for per-user agent memory.

Step 5. Recall memories across all users

Some features, such as finding similar past conversations across users whose memories the application is authorized to access, search the whole table. Use the vector index to find the approximate nearest neighbors first, and then filter. Fetch more candidates than you need, because the filter can remove some of them. However, fetching extra candidates does not guarantee that the final query returns three matches, even when more matching memories exist outside the candidate set:

SELECT * FROM ( SELECT user_id, kind, content, is_active, VEC_COSINE_DISTANCE(embedding, '[0.85, 0.15, 0.25]') AS distance FROM agent_memory ORDER BY distance LIMIT 20 ) t WHERE is_active ORDER BY distance LIMIT 3;
+---------+---------+----------------------------------------------------+-----------+-----------------------+ | user_id | kind | content | is_active | distance | +---------+---------+----------------------------------------------------+-----------+-----------------------+ | alice | fact | Alice is a data engineer at a fintech company. | 1 | 0.0040039807124229165 | | bob | fact | Bob runs a small online bookstore. | 1 | 0.0044730676562374505 | | alice | episode | Alice asked how to tune a slow join on 2026-10-01. | 1 | 0.225803083881864 | +---------+---------+----------------------------------------------------+-----------+-----------------------+

Run EXPLAIN on this query to confirm it uses the vector index: the operator info column contains annIndex: when the vector index is used. On a table with only a few rows, the optimizer might read from TiKV instead and skip the vector index; as the table grows, the optimizer might choose the vector index, depending on the query and its cost estimates. For how the placement of a filter affects the vector index, see Use the vector index with filters.

Step 6. Forget memories

To delete all memories stored for a user in the agent_memory table, for example when the user asks to be forgotten, delete by user_id:

DELETE FROM agent_memory WHERE user_id = 'bob';

Memories older than 90 days are deleted automatically by the TTL clause from Step 1. In TiDB Cloud Starter, expired rows are cleaned up by a background TTL job that runs every 15 minutes, and TTL_JOB_INTERVAL is fixed and cannot be modified. To exclude expired memories from recall at once, filter on created_at. In the Step 4 query, add AND created_at > NOW() - INTERVAL 90 DAY to the WHERE clause. In the Step 5 query, select created_at in the subquery and filter on it in the outer WHERE clause, next to is_active:

SELECT * FROM ( SELECT user_id, kind, content, is_active, created_at, VEC_COSINE_DISTANCE(embedding, '[0.85, 0.15, 0.25]') AS distance FROM agent_memory ORDER BY distance LIMIT 20 ) t WHERE is_active AND created_at > NOW() - INTERVAL 90 DAY ORDER BY distance LIMIT 3;

Do not put the created_at filter inside the subquery. A filter inside the subquery is applied before the nearest-neighbor search, so the vector index is not used.

See also

Was this page helpful?