从 SQLite+FTS5 迁移到 LanceDB:为 AI Agent 打造向量原生存储方案

  • King
  • 发布于 13小时前
  • 阅读 60

本文分享在agent-io项目中将传统SQLite全文搜索方案升级为LanceDB向量数据库的完整实践,重点介绍了如何实现语义搜索、API设计以及踩坑经验。背景在开发AIAgent项目时,我们需要一个高效的记忆存储系统来保存对话历史、用户偏好和知识库。最初选择了SQLi

本文分享在 agent-io 项目中将传统 SQLite 全文搜索方案升级为 LanceDB 向量数据库的完整实践,重点介绍了如何实现语义搜索、API 设计以及踩坑经验。

背景

在开发 AI Agent 项目时,我们需要一个高效的记忆存储系统来保存对话历史、用户偏好和知识库。最初选择了 SQLite + FTS5 的组合,因为:

  • SQLite 轻量、成熟、易于部署
  • FTS5 提供了不错的全文搜索能力

但随着 AI Agent 功能的扩展,我们发现这个方案存在明显瓶颈:

需求 SQLite+FTS5 LanceDB
关键词搜索 ✅ 支持 ✅ 支持
语义搜索 ❌ 需要额外向量库 ✅ 原生支持
相似度排序 ❌ 不支持 ✅ 支持
Schema 变更 😐 需要迁移 ✅ 灵活
部署复杂度 低(单二进制)

于是,我们决定将存储后端迁移到 LanceDB —— 一个专为 AI/ML 场景设计的嵌入式向量数据库。

为什么选择 LanceDB?

1. 向量原生设计

LanceDB 从底层开始就为向量搜索优化,使用 Apache Arrow 作为数据格式,支持高效的列式存储和向量化计算:

// 向量列定义
Field::new(
    "embedding",
    DataType::FixedSizeList(
        Arc::new(Field::new("item", DataType::Float32, true)),
        1536,  // OpenAI embedding 维度
    ),
    true,
)

2. 无服务器架构

与 Pinecone、Weaviate 等需要独立部署的服务不同,LanceDB 是嵌入式数据库:

  • 无需额外的服务进程
  • 数据存储为本地文件
  • 零运维成本

3. 统一的存储接口

一个数据库同时支持:

  • 结构化数据(元数据、时间戳等)
  • 向量数据(embeddings)
  • 全文搜索

不再需要维护多个存储系统。

技术实现

依赖配置

# Cargo.toml
[dependencies]
lancedb = "0.26"
arrow = "57"
arrow-array = "57"
arrow-schema = "57"

Schema 设计

我们的记忆条目需要存储多种类型的数据:

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

    // ... 其他字段解析
}

踩坑记录

1. Arrow 版本兼容性

问题:不同版本的 arrow crate 之间存在类型不兼容。

解决:确保所有 arrow 相关依赖使用相同版本:

arrow = "57"
arrow-array = "57"
arrow-schema = "57"

2. FixedSizeList 的 null 处理

问题:当 embedding 为 None 时,需要创建正确类型的 null 数组。

解决:使用类型标注帮助编译器推断:

FixedSizeListArray::from_iter_primitive::<Float32Type, Option<Option<f32>>, _>(
    vec![None], 1536,
)

3. RecordBatchIterator 包装

问题add 方法需要 RecordBatchReader trait,不能直接传 RecordBatch

解决:使用 RecordBatchIterator 包装:

.add(RecordBatchIterator::new(
    vec![Ok(batch.clone())], 
    batch.schema()
))

总结

从 SQLite+FTS5 迁移到 LanceDB,我们获得了:

  1. 语义搜索能力:AI Agent 可以理解用户意图,而不仅仅是匹配关键词
  2. 简化架构:不再需要维护独立的向量数据库
  3. 统一 API:一个接口处理所有查询场景
  4. 保持轻量:依然是嵌入式数据库,零运维成本

对于正在开发 AI 应用的团队,如果你的场景涉及:

  • RAG(检索增强生成)
  • 语义搜索
  • 向量存储

强烈建议尝试 LanceDB,它可能是你一直在寻找的那个「刚刚好」的解决方案。

点赞 0
收藏 0
分享
本文参与登链社区写作激励计划 ,好文好收益,欢迎正在阅读的你也加入。

0 条评论

请先 登录 后评论