数据截至 (上游 commit 7830cc746c11)
记忆与检索:schema 上下文 + 召回过往查询
30 秒导读: 一个数据 agent 要写对 SQL,得先知道"库里有哪些表和列"以及"以前类似的问题是怎么写的"。这一章讲 Wren 的
memory模块怎么低成本、按需地把这两样东西喂给 agent:schema 小就直接塞全文、schema 大就上向量检索;历史查询用本地向量库做语义召回,装不了向量库时优雅降级成纯分词 grep。schema 从哪来见 02,SQL 怎么执行见 03。
1. 这是什么(零基础也能懂)
一句话定义: wren.memory 是一个本地的、可选安装的检索层——把 MDL 语义模型(schema)和过往的"自然语言问题 → SQL"配对存进本地向量库,让 agent 在写 SQL 之前廉价地取到最相关的那几条。
解决什么问题 / 给谁用。 想象你让 AI agent 回答"上季度每个地区的营收是多少"。它要写 SQL,可你的数据库有 300 张表、几千 个字段——
- 把整个 schema 全塞进 prompt:太贵、还常常超上下文窗口。
- 只塞表名:agent 猜不出该 join 哪几张、字段叫什么。
- 让它凭空写:大概率写出语法对但语义错的 SQL(编造列名、join 错键)。
memory 模块就是来解这个"写对 SQL 前,怎么用最小代价拿到刚好够用的上下文"的问题。
它能做什么(两类记忆、两类检索):
| 记忆 | 存什么 | 检索方式 |
|---|---|---|
| Schema 记忆 | MDL 拆成的 model / column / relationship / view / cube 条目 | get_context:小库返回全文,大库向量检索 |
| 查询记忆 | 过往的 NL → SQL 配对(含种子查询) | recall_queries / MemoryIndex.search:语义相似 or grep |
用起来什么样。 它就是 wren memory 的几个子命令(装了 wren[memory] 才有语义能力):
# 把当前项目的 MDL 建成 schema 索引 + 生成种子查询
wren memory index
# 写 SQL 前,取和这个问题相关的 schema 上下文
wren memory fetch -q "revenue by region last quarter"
# 召回以前类似的 NL→SQL 例子(few-shot 用)
wren memory recall -q "revenue by region"
# 写成功后,把这条 NL→SQL 回存成新记忆
wren memory store --nl "总营收" --sql "SELECT SUM(amount) FROM orders"
一句话直觉。 把它当 agent 的外挂海马体:schema 是"这个世界长什么样"的长期记忆,查询历史是"我以前这么问过、这么答对过"的情景记忆;每次动笔前先回忆一下,别每次都从零硬想。
2. 顶层全景(它大概怎么转)
这一层由两条相互独立的检索管线组成——它们各自决定"用哪种策略",互不干扰。
怎么读下面这张图: 左边是 agent 的一次问答;中间两条竖线是两条检索管线;每条管线内部都有一个"按情况选策略"的岔口。
agent 拿到用户的自然语言问题 q
│
┌───────────────┴────────────────┐
▼ ▼
┌───────────────┐ ┌──────────────────┐
│ schema 上下文 │ │ 历史查询召回 │
│ get_context │ │ MemoryIndex │
└───────┬───────┘ └────────┬─────────┘
│ 按 schema 大小选 │ 按有没有装向量库选
┌──────┴───────┐ ┌────────┴─────────┐
▼ ▼ ▼ ▼
小库 大库 装了向量库 没装
全文塞 向量检索 LanceDB 语义 grep 分词
"full" "search" recall_queries GrepIndex
└──────┬───────┘ └────────┬─────────┘
▼ ▼
相关 schema 片段 ───────┐ ┌──── 相似 NL→SQL 例子
▼ ▼
拼进 prompt → agent 写 SQL → 见 03 执行
│
▼(成功后)
store_query 回存成新记忆
部件一句话职责:
| 部件 | 干什么 | 文件 |
|---|---|---|
WrenMemory | 面向程序的高层 API 门面 | memory/__init__.py:17 |
MemoryStore | LanceDB 两张表的读写与所有策略逻辑 | memory/store.py:74 |
MemoryIndex / GrepIndex / LanceDBIndex | 可插拔的查询召回后端(语义 or grep) | memory/index_backend.py:47 |
schema_indexer | 把 MDL manifest 拆成可检索的 schema 条目 + 全文描述 | memory/schema_indexer.py:309 |
seed_queries | 从 manifest 自动造几条 canonical NL→SQL 种子 | memory/seed_queries.py:31 |
markdown | knowledge/sql/*.md 的读写(查询记忆的真源) | memory/markdown.py:69 |
embeddings | 本地 sentence-transformers 嵌入函数 | memory/embeddings.py:77 |
主线走一遍(高层): 用户问题进来 → 左管线取 schema、右管线取历史例子 → 两者拼进 prompt → agent 写出模型名 SQL(交给 03 治理执行)→ 成功的配对回存,喂养下一次召回。
3. 核心原理(逐个机制,由浅入深)
3.1 自适应 schema 上下文:小库塞全文,大库才检索
它要解决的小问题。 "取相关 schema"听起来天然就是"向量检索最相关的几条"。但对小 schema,检索反而是坏主意——你把它切成碎片、只捞回 5 条,agent 就看不到表之间的 join 路径和整体结构了。
思路/直觉。 于是 Wren 做了个很朴素但关键的判断:
如果整个 schema 的纯文本描述短到能舒服地塞进一个上下文窗口,就整个塞进去(策略
full),让 LLM 看到完整结构;只有当它太大时,才退而求其次做向量检索(策略search)。
阈值是字符数,常量 SCHEMA_DESCRIBE_THRESHOLD = 30_000(约 8K token)。
为什么用字符数而不是 token 数? 源码注释讲得很直白:文本生成完后数字符是免费的,而精确数 token 要拉一个 tokenizer 依赖;30K 字符 ≈ 8K token 是按英文 4:1 估的,CJK 约 1.5:1,所以中文密集的 schema 会更早切到检索——这是偏安全的方向。见 schema_indexer.py:60-70 常量与注释。
原理演示(示意,非源码):
# 演示 get_context 的岔口:先算全文,短就全给,长才检索
def get_context(manifest, query, threshold=30_000):
text = describe_schema(manifest) # 把 MDL 摊成结构化纯文本
if len(text) <= threshold: # 小库
return {"strategy": "full", "schema": text}
# 大库:嵌入 query,在 schema_items 表里做向量近邻
results = search_schema(query, mdl_hash=hash(manifest))
return {"strategy": "search", "results": results}
# 重点看:走不走检索,由 schema 的“全文长度”决定,而不是一律检索
真实实现。 岔口在 store.py:300-312 的 get_context:describe_schema(manifest) 生成全文,len(text) <= threshold 就返回 {"strategy":"full", ...},否则调 _search_schema 返回 {"strategy":"search", ...}。全文生成器 describe_schema 在 schema_indexer.py:94,按 model → column → relationship → view → cube 逐层排成缩进文本。
关键细节。 向量检索这条(_search_schema,store.py:314-345)会始终用当前 manifest 的哈希 mdl_hash 作过滤条件(store.py:334),外加可选的 item_type / model_name 过滤——保证只捞回属于当前 schema 版本的条目,不会串进旧 schema 的残留行。
3.2 把 manifest 拆成"可检索的 schema 条目"
它要解决的小问题。 向量检索的前提是"有一堆可以逐条嵌入、逐条命中的东西"。MDL manifest 是一棵嵌套的 JSON 树,得先摊平成一条条独立记录。
思路。 extract_schema_items(schema_indexer.py:309)遍历 manifest,给每一个 model、每一列 column、 每条 relationship、每个 view,以及 cube 及其 measure / dimension / time_dimension,各生成一条记录。每条记录有一个给人也给嵌入模型看的 text 字段(合成的自然语言描述),外加结构化元数据列。
举例,一列会被合成成这样一句可嵌入的话(见 _column_record,schema_indexer.py:404):
Column 'amount' (double) in model 'orders': 订单金额. Constraints: NOT NULL.
版本指纹。 每条记录都盖上 mdl_hash——manifest 的 SHA-256 前 16 位(manifest_hash,schema_indexer.py:14)。它有意排除 _ 开头的内部键(如 _instructions),这样"只改了给 agent 的说明、没动 schema"时缓存不失效。schema_is_current(store.py:258)就靠"表里每一行的 mdl_hash 是否都等于当前哈希"来判断索引是否新鲜。
3.3 可插拔召回后端:语义优先,grep 兜底(hybrid / 优雅降级)
它要解决的小问题。 语义召回要装 lancedb + sentence-transformers(几百 MB),不是人人都愿意装。但"召回相似历史查询"这件事不该因为没装重依赖就彻底用不了。
思路。 把召回抽象成一个接口 MemoryIndex(index_backend.py:47),给两个实现,能力不同但接口一致:
| 后端 | 依赖 | 怎么召回 | 索引在哪 |
|---|---|---|---|
LanceDBIndex | lancedb + sentence-transformers | 向量语义近邻 | LanceDB(派生索引) |
GrepIndex | 无(纯标准库) | 分词重叠 + 子串打分 | 无——knowledge/sql/ 本身就是索引 |
这就是 README 说的 "hybrid retrieval":有重依赖时走语义、没有时走 grep,两条路产出同样形状的结果行(靠 _pair_to_result,index_backend.py:32 对齐字段)。
GrepIndex 的打分很朴素但够用(index_backend.py:90-106):
# 示意,非源码:grep 后端怎么给一条历史配对打分
q_tokens = tokens(query) # query 分词(长度≥2 的字母数字)
score = len(q_tokens & (tokens(pair.nl) | tokens(pair.sql))) # 词重叠个数
if query.lower() in pair.nl.lower():
score += 5 # 整句子串命中,直接顶到最前
# 按 score 降序;同分时用 nl 字典序稳定排序(结果可复现)
真实分词器 _tokens(index_backend.py:28)只取长度 ≥2 的 [a-z0-9]+;整句子串命中额外 +5(index_backend.py:101),让"问得几乎一模一样"的历史排最前。
降级是自动且可预测的。 resolve_backend(index_backend.py:147)决定实际用哪个:环境变量 WREN_MEMORY_BACKEND=grep 强制 grep;否则当且仅当 lancedb 与 sentence_transformers 两个包都可导入(_extra_available,index_backend.py:143)才用 lancedb,否则降级 grep。get_index(index_backend.py:166)据此构造实例——即使显式传 backend="lancedb" 但没装依赖,也会安静地退回 GrepIndex。
3.4 查询记忆:存、召回、去重、导入导出
存(store_query)。 一条 NL→SQL 配对进来,store_query(store.py:349)只对 NL 部分做嵌入(因为召回是"拿新问题去匹配老问题"),连同 SQL、datasource、tags、时间戳写进 query_history 表。
召回(recall_queries)。 recall_queries(store.py:382)把新问题嵌入,在表里做向量近邻,可选按 datasource 过滤,返回前 limit 条(默认 3),并丢掉 vector 列只回文本。
导入 / 导出(queries.yml)。 dump_queries(store.py:485)把整张查询表导成不含向量的字典列表(CLI 再写成 queries.yml);load_queries(store.py:565)反向批量导入,支持三种模式:
| 模式 | 行为 |
|---|---|
| 默认(skip) | 已存在的 (nl, sql) 精确重复直接跳过 |
upsert=True | 按 nl 归并,同 NL 的旧行全删再插新(NL 相同视作同一逻辑配对) |
overwrite=True | 先按 source 清空同源全部,再全量插入 |
去重靠一张预建索引。 _existing_pairs_index(store.py:500)一次性扫全表,建两个查找结构:exact_set = {(nl, sql)} 供 skip 模式做精确判重;nl_to_rowids = {nl: [行号...]} 供 upsert 模式定位"同一个问题的所有旧行"并批量删除。这样导入不会 O(n²) 反复查表。
列表与遗忘。 list_queries(store.py:408)分页列出并给每行附上 _row_id(未过滤表中的原始位置索引),forget_queries_by_ids(store.py:449-470)按这些位置删除,forget_queries_by_source(store.py:473)按 source:* tag 批量删——_row_id 用原始索引正是为了"即便加了 source 过滤,删的仍是正确的行"。
3.5 种子查询:冷启动就有例子可召回
它要解决的小问题。 一个刚建索引的新项目,查询历史是空的,recall 什么都召不回,few-shot 没料。
思路。 建 schema 索引时顺手从 manifest 造几条 canonical 的 NL→SQL 灌进去当"启动记忆"。_upsert_seed_queries(store.py:234)先把新种子全部算好、嵌完向量,再做原子替换——删掉带 source:seed tag 的旧种子、批量写入新的(用户确认过的条目不受影响)——避免"旧的删了、新的没写成"的空窗;种子由 generate_seed_queries(seed_queries.py:31)生成。
造哪些(模板化,seed_queries.py:57-158): 每个非 raw 层的 model 造"列全部 / 求和 / 分组求和";有 acceptedValues 的列造"按枚举值过滤";每条 relationship 造一条 join 例子。
一个巧妙的细节:不给 join 键求和。 它会挑"第一个数值列"来做 SUM,但跳过主键、关系键、以及名字像 *_id 的列(seed_queries.py:84-96 的 is_identifier 判断)——因为 SUM(customer_id) 语义上是废话。关系键从哪来?解析每条 relationship 的 condition(用 sqlglot)抽出两侧列名(_relationship_key_columns,seed_queries.py:161)。
4. 深入实现:两张 Arrow 表与维度对齐
查询记忆和 schema 记忆各存一张 LanceDB 表,表结构用 PyArrow schema 显式声明(不靠推断)。
schema_items 表(_schema_items_arrow_schema,store.py:35)每行是一个 schema 条目:
| 列 | 类型 | 作用 |
|---|---|---|
text | utf8 | 合成的自然语言描述(被嵌入的对象) |
vector | list<float32>[dim] | 嵌入向量,长度 = 模型维度 |
item_type / model_name / item_name | utf8 | 过滤维度(model/column/...) |
data_type / expression / is_calculated | utf8 / bool | 列的类型、计算表达式 |
mdl_hash | utf8 | schema 版本指纹,检索时强制匹配 |
indexed_at | timestamp(UTC) | 建索引时间 |
query_history 表(_query_history_arrow_schema,store.py:52)每行是一条 NL→SQL:text/vector(NL 的嵌入)、nl_query、sql_query、datasource、created_at、tags。
维度必须对齐。 向量列的长度 dim 不能写死,得等于嵌入模型真实输出的维度。embeddings.py 负责这件事:
- 默认模型
_DEFAULT_MODEL = paraphrase-multilingual-MiniLM-L12-v2(可用WREN_EMBEDDING_MODEL覆盖),默认维度_DEFAULT_DIM = 384(embeddings.py:13-16)。多语言模型 = 中文英文问题都能嵌。 get_embedding_function(embeddings.py:77)返回一个 local-first 的 sentence-transformers 嵌入函数(LanceDBSentenceTransformerEmbeddings的子类适配器):先从本地 HF 缓存加载、 失败才回退联网加载,模型构造每进程单飞(embeddings.py:27-74)。本地跑、无需 API key。warm_up(embeddings.py:109)嵌一个"probe"触发模型加载并返回真实维度。维度解析现在 是惰性的:MemoryStore.__init__不再急切加载模型,_embed_fn/_dim是带锁的属性 (store.py:96-120),首次用到才初始化;_resolve_dim(store.py:133)优先沿用已有表的向量维度,每次写入前_validate_and_set_dim对账——换了不同维度的WREN_EMBEDDING_MODEL去写旧库会当场报清晰的ValueError(提示重置 store 或换回原模型),不许静默混维度。加载时的原生进度条/噪声用suppress_stderr静音(embeddings.py:91)。
5. 一段白话流程:agent 检索 → 写 SQL → 回存
把前面拼起来,一次完整问答里 memory 的位置:
1. 用户问:"上季度每个地区的营收?"
2. agent 调 get_context(manifest, q)
· schema 全文 ≤ 30K 字符 → 直接拿到整份结构化 schema(策略 full)
· schema 很大 → 向量检索,拿回最相关的 model/column 几条(策略 search)
3. agent 调 recall / MemoryIndex.search(q)
· 装了向量库 → LanceDB 语义召回 2~3 条相似的老 NL→SQL 当 few-shot
· 没装 → GrepIndex 分词打分召回(能力弱些,但不空手)
4. agent 把 [schema 上下文] + [few-shot 例子] 拼进 prompt,写出模型名 SQL
5. SQL 交给语义引擎治理、执行(见 03);成功拿到结果
6. agent 调 store_query(nl, sql):
· 先写 knowledge/sql/<slug>.md(真源,dependency-free)
· 若装了向量库,再嵌入 NL 入 query_history 表(派生索引)
→ 这条经验从此可被 recall 召回,喂养下一次
记住"真源 vs 派生"这条线。 查询记忆的权威来源是 knowledge/sql/*.md(markdown.py:69 的 load_query_pairs),LanceDB 只是从它派生出来的、可随时重建的索引——正如 wren context build 把 YAML 编译成 target/mdl.json。GrepIndex 干脆直接读这些 markdown、不建任何派生物(index_backend.py:80-88)。queries.yml 则是另一条批量导入导出通道,直接作用在 LanceDB 存储上(CLI 里还保留它作 legacy 兼容)。