📣
TiDB Cloud Premium 开放公测中。为企业级工作负载提供无限扩展、即时弹性伸缩和高级安全保障。此页面由 AI 自动翻译,英文原文请见此处。
​
​

在 TiDB 中存储 AI Agent Memory



一个能够跨会话记住用户的 AI agent,其 memory store 需要完成四件事:按语义查找记忆、按用户和类型缩小范围、在不留下半写入状态的情况下替换过时记忆,以及在记忆过期或用户提出要求时将其遗忘。在 TiDB 中,一张表和普通 SQL 就能覆盖这四种需求,因为向量搜索、二级索引、事务和行过期都位于同一个数据库中。

本文将逐步介绍这张表的设计。若要查看基于相同思路、使用 pytidb Python SDK 构建的完整聊天应用,请参见AI Agent Memory 示例。

前提条件

  • 一个 TiDB Cloud Starter 实例。你可以在 TiDB Cloud 上创建一个免费的 TiDB Cloud Starter 实例。
  • 一个已连接到该实例的 MySQL 客户端。在 TiDB Cloud console 中,进入 TiDB Cloud Starter 实例的概览页面,然后点击 Connect 获取连接参数。

本文中的这些示例使用了 3 维向量,并手动指定其值,从而使 SQL 示例和查询结果更加容易理解。这些向量仅用于演示,不代表由嵌入模型生成的向量嵌入。在你的应用中,请将向量维度设置为嵌入模型的输出大小,例如 VECTOR(1536),并使用该模型生成每个 embedding。

第 1 步:创建 memory 表

首先选择一个数据库。本示例使用每个 TiDB Cloud Starter 实例都自带的 test 数据库:

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';

表中的每个部分都对应一种记忆操作:

  • embedding 和向量索引 idx_embedding 用于按语义查找记忆。
  • user_id 和 kind 用于将记忆限定到某个用户和某种类型。建立在 (user_id, is_active) 上的二级索引 idx_user 可以加快读取某个用户的活跃记忆。
  • is_active 用于让过时记忆退役,而无需删除它。
  • TTL 子句使 TiDB 自动清理已存储超过 90 天的记忆。过期记忆不一定会立即删除,因为 TTL 清理是周期性运行的。更多信息,请参见 TTL (Time to Live)。

第 2 步:写入记忆

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]');

第 3 步:在事务中替换过时记忆

当用户修改偏好时,在一个事务中让旧记忆退役并写入新记忆。事务可以确保读取方不会看到只完成一部分的替换结果。

为防止并发写入为同一用户创建多个活跃偏好,请按为每个用户最多保留一个活跃偏好中的说明添加唯一性约束。

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;

检查结果:

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 | +----------------------------------+-----------+

如果事务中的某条语句失败,例如某个值不符合其列定义,请执行 ROLLBACK 而不是 COMMIT。语句失败不会结束事务:已经成功的语句仍处于待提交状态,而如果执行 COMMIT,这些修改仍会被保存,例如导致用户的旧偏好已退役但新偏好没有写入。执行 ROLLBACK 后,记忆会保持原样。

为每个用户最多保留一个活跃偏好

事务可以保证每次替换的原子性,但并发写入仍可能为同一用户创建多个活跃偏好。例如,当当前不存在活跃偏好时,两个事务都可能各自更新 0 行、插入一条新偏好,并成功提交。

如果你的应用只允许每个用户有一个活跃偏好,请添加以下生成列和唯一索引。对于已有表,请先为每个用户停用或移除重复的活跃偏好。

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);

该生成列仅在活跃偏好行中包含 user_id。唯一索引可防止两行此类记录拥有相同的用户 ID。其他行在该列中为 NULL,因此该索引允许每个用户存在多个非活跃偏好、fact 和 episode。

添加唯一索引后,按上文所示,在单个事务中停用旧偏好并插入新偏好。即使在并发写入场景下,唯一索引也能强制执行这一限制。如果替换失败,请回滚事务。若要重试,请开启一个新事务,并重新执行 UPDATE 和 INSERT。

第 4 步:召回单个用户的记忆

大多数 agent 请求只需要当前用户的记忆。先按用户过滤,再按与当前请求 embedding 的距离进行排序:

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 | +------------+----------------------------------------------------+-----------------------+

该查询不会使用向量索引,因为 WHERE 过滤会在最近邻搜索之前应用。它可以使用二级索引 idx_user,只读取该用户的活跃记忆,并为每条记录计算精确距离。结果是精确的;在使用二级索引时,其成本主要取决于匹配记忆的数量,而不需要全表扫描。对于按用户维度的 agent memory,这通常是正确的模式。

第 5 步:跨所有用户召回记忆

某些功能需要搜索整张表,例如在应用已被授予权限可访问的记忆范围内,查找不同用户之间相似的历史对话。此时应先使用向量索引查找近似最近邻,再进行过滤。获取的候选项数量应多于最终所需数量,因为过滤可能会移除其中一部分。不过,即使候选集之外还存在更多匹配记忆,获取更多候选项也不能保证最终查询一定返回三条匹配结果:

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 | +---------+---------+----------------------------------------------------+-----------+-----------------------+

对该查询执行 EXPLAIN 以确认其使用了向量索引:当使用向量索引时,operator info 列会包含 annIndex:。如果表中只有少量行,优化器可能会改为从 TiKV 读取并跳过向量索引;随着表变大,优化器可能会根据查询及其成本估算选择向量索引。关于过滤条件的位置如何影响向量索引,请参见在带过滤条件的查询中使用向量索引。

第 6 步:遗忘记忆

要删除 agent_memory 表中为某个用户存储的所有记忆,例如当用户要求被遗忘时,可按 user_id 删除:

DELETE FROM agent_memory WHERE user_id = 'bob';

超过 90 天的记忆会由第 1 步中的 TTL 子句自动删除。在 TiDB Cloud Starter 中,过期行由后台 TTL 任务每 15 分钟清理一次,且 TTL_JOB_INTERVAL 是固定值,无法修改。若要立即在召回结果中排除过期记忆,请基于 created_at 进行过滤。在第 4 步的查询中,将 AND created_at > NOW() - INTERVAL 90 DAY 添加到 WHERE 子句中。在第 5 步的查询中,在子查询里选出 created_at,并在外层 WHERE 子句中与 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;

不要将 created_at 过滤条件放在子查询内部。子查询内部的过滤会在最近邻搜索之前应用,因此不会使用向量索引。

另请参阅

文档内容是否有帮助?