本文分享在agent-io项目中将传统SQLite全文搜索方案升级为LanceDB向量数据库的完整实践,重点介绍了如何实现语义搜索、API设计以及踩坑经验。背景在开发AIAgent项目时,我们需要一个高效的记忆存储系统来保存对话历史、用户偏好和知识库。最初选择了SQLi
本文分享在 agent-io 项目中将传统 SQLite 全文搜索方案升级为 LanceDB 向量数据库的完整实践,重点介绍了如何实现语义搜索、API 设计以及踩坑经验。
在开发 AI Agent 项目时,我们需要一个高效的记忆存储系统来保存对话历史、用户偏好和知识库。最初选择了 SQLite + FTS5 的组合,因为:
但随着 AI Agent 功能的扩展,我们发现这个方案存在明显瓶颈:
| 需求 | SQLite+FTS5 | LanceDB |
|---|---|---|
| 关键词搜索 | ✅ 支持 | ✅ 支持 |
| 语义搜索 | ❌ 需要额外向量库 | ✅ 原生支持 |
| 相似度排序 | ❌ 不支持 | ✅ 支持 |
| Schema 变更 | 😐 需要迁移 | ✅ 灵活 |
| 部署复杂度 | 低 | 低(单二进制) |
于是,我们决定将存储后端迁移到 LanceDB —— 一个专为 AI/ML 场景设计的嵌入式向量数据库。
LanceDB 从底层开始就为向量搜索优化,使用 Apache Arrow 作为数据格式,支持高效的列式存储和向量化计算:
// 向量列定义
Field::new(
"embedding",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
1536, // OpenAI embedding 维度
),
true,
)
与 Pinecone、Weaviate 等需要独立部署的服务不同,LanceDB 是嵌入式数据库:
一个数据库同时支持:
不再需要维护多个存储系统。
# Cargo.toml
[dependencies]
lancedb = "0.26"
arrow = "57"
arrow-array = "57"
arrow-schema = "57"
我们的记忆条目需要存储多种类型的数据:
fn schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("id", DataType::Utf8, false),
Field::new("content", DataType::Utf8, false),
Field::new("embedding", DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
1536,
), true),
Field::new("memory_type", DataType::Utf8, false),
Field::new("metadata", DataType::Utf8, true),
Field::new("created_at", DataType::Int64, false),
Field::new("last_accessed", DataType::Int64, true),
Field::new("importance", DataType::Float32, false),
Field::new("access_count", DataType::UInt32, false),
]))
}
LanceDB 支持内存模式和文件持久化两种模式:
pub struct LanceDbStore {
table: Arc<Table>,
}
impl LanceDbStore {
// 内存模式 - 适合测试和临时存储
pub async fn new() -> Result<Self> {
Self::open_uri("memory://agent_io_memories").await
}
// 文件模式 - 适合生产环境持久化
pub async fn open<P: Into<PathBuf>>(path: P) -> Result<Self> {
let path = path.into();
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
Self::open_uri(&path.to_string_lossy()).await
}
async fn open_uri(uri: &str) -> Result<Self> {
let db = connect(uri).execute().await?;
let table = if db.table_names().execute().await?
.contains(&"memories".to_string())
{
db.open_table("memories").execute().await?
} else {
db.create_empty_table("memories", Self::schema())
.execute().await?
};
Ok(Self { table: Arc::new(table) })
}
}
LanceDB 使用 Arrow RecordBatch 作为数据格式,需要将结构体转换为列式存储:
async fn add(&self, entry: MemoryEntry) -> Result<String> {
let id = entry.id.clone();
// 构建各个列
let id_array = StringArray::from(vec![entry.id]);
let content_array = StringArray::from(vec![entry.content]);
// 向量列需要特殊处理
let embedding_array = if let Some(ref embedding) = entry.embedding {
FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
vec![Some(embedding.iter().map(|&v| Some(v)).collect())],
1536,
)
} else {
FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
vec![None], 1536,
)
};
// 组装 RecordBatch
let batch = RecordBatch::try_new(schema, vec![
Arc::new(id_array),
Arc::new(content_array),
Arc::new(embedding_array),
// ... 其他列
])?;
// 写入数据
self.table
.add(RecordBatchIterator::new(
vec![Ok(batch.clone())],
batch.schema()
))
.execute()
.await?;
Ok(id)
}
对于不需要语义理解的场景,传统的关键词搜索仍然有用:
async fn search(&self, query: &str, limit: usize) -> Result<Vec<MemoryEntry>> {
let batches = self.table
.query()
.only_if(format!("content LIKE '%{}%'", query))
.limit(limit)
.execute()
.await?
.try_collect::<Vec<_>>()
.await?;
// 解析结果...
}
这是迁移到 LanceDB 的主要收益 —— 支持基于向量相似度的语义搜索:
async fn search_by_embedding(
&self,
embedding: &[f32],
limit: usize,
threshold: f32,
) -> Result<Vec<MemoryEntry>> {
let batches = self.table
.query()
.nearest_to(embedding) // 核心API:向量搜索
.limit(limit * 2)
.execute()
.await?
.try_collect::<Vec<_>>()
.await?;
let mut results: Vec<(MemoryEntry, f32)> = Vec::new();
for batch in batches {
for i in 0..batch.num_rows() {
let entry = parse_batch_row(&batch, i)?;
// LanceDB 返回 _distance 列表示距离
let similarity = batch.column_by_name("_distance")
.and_then(|col| col.as_any().downcast_ref::<Float32Array>())
.map(|arr| 1.0 - arr.value(i)) // 距离转相似度
.unwrap_or(0.0);
if similarity >= threshold {
results.push((entry, similarity));
}
}
}
// 按相似度排序
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Equal));
results.truncate(limit);
Ok(results.into_iter().map(|(e, _)| e).collect())
}
从 RecordBatch 中提取数据需要类型转换:
fn parse_batch_row(batch: &RecordBatch, i: usize) -> Result<MemoryEntry> {
// 普通字符串列
let id = batch.column(0)
.as_any()
.downcast_ref::<StringArray>()
.map(|arr| arr.value(i).to_string())
.unwrap_or_default();
// 向量列需要特殊处理
let embedding = batch.column(2)
.as_any()
.downcast_ref::<FixedSizeListArray>()
.and_then(|arr| {
if arr.is_null(i) { return None; }
arr.value(i)
.as_any()
.downcast_ref::<Float32Array>()
.map(|v| v.values().to_vec())
});
// ... 其他字段解析
}
问题:不同版本的 arrow crate 之间存在类型不兼容。
解决:确保所有 arrow 相关依赖使用相同版本:
arrow = "57"
arrow-array = "57"
arrow-schema = "57"
问题:当 embedding 为 None 时,需要创建正确类型的 null 数组。
解决:使用类型标注帮助编译器推断:
FixedSizeListArray::from_iter_primitive::<Float32Type, Option<Option<f32>>, _>(
vec![None], 1536,
)
问题:add 方法需要 RecordBatchReader trait,不能直接传 RecordBatch。
解决:使用 RecordBatchIterator 包装:
.add(RecordBatchIterator::new(
vec![Ok(batch.clone())],
batch.schema()
))
从 SQLite+FTS5 迁移到 LanceDB,我们获得了:
对于正在开发 AI 应用的团队,如果你的场景涉及:
强烈建议尝试 LanceDB,它可能是你一直在寻找的那个「刚刚好」的解决方案。
如果觉得我的文章对您有用,请随意打赏。你的支持将鼓励我继续创作!