PandasAI LanceDB 向量存储扩展：为 Agent 训练与 RAG 检索提供持久化向量能力-创锋一号

PandasAI LanceDB 向量存储扩展：为 Agent 训练与 RAG 检索提供持久化向量能力

【免费下载链接】pandas-aiChat with your database or your datalake (SQL, CSV, parquet). PandasAI makes data analysis conversational using LLMs and RAG.项目地址: https://gitcode.com/GitHub_Trending/pa/pandas-ai

本文面向需要在 PandasAI 中使用向量数据库为 Agent 提供"训练记忆"的开发者。通过 pandasai-lancedb 扩展，你可以将问题-代码对（Q&A）与业务文档持久化到本地 LanceDB，并在对话时按语义相似度检索，作为上下文喂给大模型。读完本文，你将掌握该扩展的安装方式、LanceDB类的全部核心 API、默认/自定义嵌入模型的选择，以及它如何与Agent.train()无缝衔接形成完整的 RAG 闭环。

PandasAI 通过VectorStore抽象层（pandasai/vectorstores/vectorstore.py）为 Agent 提供"训练数据"的存取能力：把历史问答（自然语言问题 + 对应生成的 Python 代码）和补充文档向量化后持久化，在后续会话中按语义相似度召回，作为上下文注入提示词。pandasai-lancedb正是该抽象层基于 LanceDB 的官方实现，无需额外部署服务端，数据以本地文件形式持久化，开箱即用。

一、扩展概览与安装

1.1 它解决什么问题

该扩展将 LanceDB 与 PandasAI 集成，提供面向"增强数据分析与机器学习任务"的向量存储能力（见 README.md）。从源码结构看，它通过LanceDB类实现 VectorStore 抽象基类定义的整套接口，包括：

问题-代码对的增删改查与语义检索（add_question_answer/update_question_answer/delete_question_and_answers/get_relevant_question_answers）；
文档的增删改查与语义检索（add_docs/update_docs/delete_docs/get_relevant_docs）；
按 ID 精确获取记录、按相似度距离阈值过滤结果等进阶能力。

1.2 安装方式

官方 README 推荐的安装方式是通过 Poetry：

poetry add pandasai-lancedb

对应包元信息见 pyproject.toml，核心依赖与版本约束如下：

依赖项	版本约束	说明
`python`	`>=3.8,<3.12`	注意：不支持 Python 3.12+
`pandasai`	`>=3.0.0b4`	依赖主框架的`VectorStore`抽象接口
`lancedb`	`^0.5.0`	LanceDB 客户端库
`numpy`	`1.23.2`	锁定版本
`sentence-transformers`	`^2.2.2`	默认嵌入模型加载器

测试依赖（可选分组test）包含pytest ^7.4.0、pytest-cov、pytest-mock，用于运行 test_lancedb.py 中的单元测试。

提示：该包位于extensions/ee/目录下，属于 Enterprise（企业版）扩展。包的 LICENSE 文件（LICENSE）声明为 Sinaptik GmbH Enterprise License，商业使用需联系pm@sinaptik.ai（见 README.md）。

二、核心类`LanceDB`与构造参数详解

扩展的唯一对外入口是LanceDB类（由 pandasai_lancedb/init.py 导出），实现位于 pandasai_lancedb/lancedb.py。

2.1 构造参数

from pandasai_lancedb import LanceDB vector_store = LanceDB( table_name="pandasai", # 表名前缀，默认 "pandasai" embedding_function=None, # 自定义嵌入函数，默认 None persist_path="/tmp/lancedb", # 数据持久化目录，默认 /tmp/lancedb max_samples=1, # 检索返回的最大样本数，默认 1 similary_threshold=1.5, # 相似度距离阈值，默认 1.5 logger=None, # 日志器，默认创建新 Logger )

参数	默认值	作用与注意点
`table_name`	`"pandasai"`	实际会创建两张表：`{table_name}-qa`与`{table_name}-docs`
`embedding_function`	`None`	自定义嵌入函数，签名`Callable[[List[str]], List[float]]`；为`None`时使用默认嵌入模型
`persist_path`	`"/tmp/lancedb"`	LanceDB 数据持久化目录，`lancedb.connect()`直接连接该路径
`max_samples`	`1`	检索时的默认 top-k 数量
`similary_threshold`	`1.5`	距离阈值，`_distance < threshold`的结果才会被保留
`logger`	`None`	传入`pandasai.helpers.logger.Logger`实例用于日志输出

注意：源码中该参数名为similary_threshold（保留原始拼写），这是当前仓库的实际 API 签名，使用时请照此书写。

2.2 初始化流程（源码级解读）

构造时LanceDB会依次完成以下步骤（见 lancedb.py）：

调用lancedb.connect(persist_path)连接（必要时创建）本地数据库；
根据是否传入embedding_function构造Schema并创建两张 LanceDB 表；
表已存在则open_table打开，不存在则create_table，保证重复初始化不丢数据；
通过 Logger 输出持久化路径与初始化成功的日志。

从源码可以推断，Schema类（同文件 lancedb.py）会定义两个继承自lancedb.pydantic.LanceModel的模型：

QA_pairs：字段为id（str）、qa（嵌入源字段SourceField）、metadata（str）、vector（Vector(ndims)向量字段）；
Docs：字段为id、doc（嵌入源字段）、metadata、vector。

两者共享同一个嵌入函数，向量的维度由嵌入模型动态计算（ndims()会先对["foo"]生成一次嵌入以确定维度）。

三、嵌入模型：默认模型与自定义函数

3.1 默认嵌入方案

当embedding_function=None时，Schema使用 LanceDB 注册表中的sentence-transformers嵌入函数：

get_registry().get("sentence-transformers").create( name="BAAI/bge-small-en-v1.5", device="cpu" )

即默认使用BAAI/bge-small-en-v1.5模型并在 CPU 上推理（见 lancedb.py）。该默认行为意味着首次使用时需要下载模型权重。

另外，get_embeddings()方法中单独使用了BAAI/bge-large-zh-v1.5（见 lancedb.py），说明该实现同时关注中英文语义场景，从源码结构看它主要用于直接获取文本向量的辅助场景。

3.2 自定义嵌入函数

若希望使用自有嵌入模型，可在构造时传入一个可调用对象，例如：

from sentence_transformers import SentenceTransformer model = SentenceTransformer("BAAI/bge-small-en-v1.5") def embed(texts): return model.encode(texts, normalize_embeddings=True).tolist() vector_store = LanceDB(embedding_function=embed)

传入后，Schema会通过EmbeddingFunctionRegistry.get_instance(model)注册并创建名为embedding_function的嵌入器（见 lancedb.py）。同时，所有"写入"（add_*）与"检索"（get_relevant_*）路径都会改为调用你的函数生成向量（见 lancedb.py），保证写入与查询使用同一套向量空间。

四、数据写入：Q&A 对与文档的持久化

4.1 写入 Q&A 对

ids = vector_store.add_question_answer( queries=["What is LanceDB?", "How does it work?"], codes=["print('Hello')", "for i in range(10): print(i)"], # ids=["custom_id_1", "custom_id_2"], # 可选，默认自动生成 UUID # metadatas=[{"source": "docs"}, ...], # 可选元数据 )

实现要点（见 lancedb.py）：

queries与codes长度不一致时会抛出ValueError并提示维度不匹配；
未传入ids时自动生成形如{uuid4}-qa的 ID；
未传入metadatas时统一填充字符串"None"；
每条记录按Q: {query}\n A: {code}的格式拼接为qa字段（格式定义在 VectorStore._format_qa）；
有自定义嵌入函数时，预先计算vector列一并写入；否则交由 LanceDB 的嵌入源字段在写入时自动生成。

4.2 写入文档

ids = vector_store.add_docs( docs=["Document 1", "Document 2"], # ids=["doc_id_1", "doc_id_2"], # 可选 # metadatas=[{"topic": "intro"}, ...], # 可选 )

逻辑与 Q&A 写入对称：ID 默认形如{uuid4}-docs，元数据默认"None"，有自定义嵌入函数时预计算向量（见 lancedb.py）。

4.3 更新与删除

update_question_answer(ids, queries, codes, metadatas)：按id = '{id}'条件对qa表执行 LanceDB 行更新；同样会校验 queries/codes 长度一致（见 lancedb.py）；
update_docs(ids, docs, metadatas)：按 ID 更新doc与metadata字段（见 lancedb.py）；
delete_question_and_answers(ids)/delete_docs(ids)：按 ID 逐条执行table.delete(f"id = '{id}'")，成功返回True（见 lancedb.py）。

五、语义检索：距离阈值与 Top-K

检索是本扩展服务于 RAG 的核心能力：

# 检索与问题最相关的 Q&A 对 results = vector_store.get_relevant_question_answers("What is LanceDB?", k=2) # 检索相关文档 doc_results = vector_store.get_relevant_docs("How does it work?", k=3)

两者的返回结构一致（见 lancedb.py）：

{ "documents": [["Q: What is LanceDB?\n A: print('Hello')", ...]], "metadatas": [["None", ...]], }

实现要点：

k未指定时使用构造参数max_samples；
默认嵌入模式下直接以问题文本调用table.search(query=question)；自定义嵌入模式下先用embedding_function([question])生成查询向量再搜索；
结果经过_filter_docs_based_on_distance过滤：仅保留_distance < similary_threshold（默认 1.5）的记录，再拆分为documents与metadatas两个列表，与VectorStore接口约定的返回格式对齐。

此外还提供了按 ID 精确获取的方法：get_relevant_question_answers_by_id(ids)与get_relevant_docs_by_id(ids)，通过.where(f"id = '{id}'")过滤并仅选择metadata与正文列（见 lancedb.py）。

六、与 Agent 的集成：train() 训练闭环

该扩展真正的使用场景是配合 PandasAI Agent 的train()方法。在 pandasai/agent/base.py 中，train()的逻辑如下：

若 Agent 未配置vectorstore，抛出MissingVectorStoreError；
若只传了queries或codes之一，抛出ValueError；
docs非空时调用vectorstore.add_docs(docs)；
queries与codes同时提供时调用vectorstore.add_question_answer(queries, codes)。

组合使用示例：

from pandasai import Agent from pandasai_lancedb import LanceDB vector_store = LanceDB( table_name="my_agent", persist_path="/data/lancedb", ) agent = Agent(dataframes, vectorstore=vector_store) # 训练：注入历史问答与补充文档 agent.train( queries=["What is LanceDB?", "How to filter a dataframe?"], codes=["print('Hello')", "df[df['col'] > 0]"], docs=["This is a supplementary document about the dataset."], )

训练完成后，对话阶段 Agent 会从向量库召回相关知识并注入提示词上下文（仓库中 vectordb_docs.tmpl 即用于承载这部分召回内容，可从其命名与位置推断其用途）。

七、单元测试与验证

扩展自带完整单元测试（test_lancedb.py），可验证上述全部行为：

构造默认参数与自定义 Logger；
自动创建pandasai-qa/pandasai-docs两张表（表不存在时创建）；
add_question_answer/add_docs支持自动 ID 与自定义 ID；
输入维度不匹配时抛出ValueError；
update_question_answer/delete_question_and_answers/delete_docs的基本行为；
get_relevant_question_answers/get_relevant_docs的返回结构与按 ID 查询。

测试在tearDown中清理/tmp/lancedb目录，因此可以在本机直接运行：

pytest extensions/ee/vectorstores/lancedb/tests/test_lancedb.py

八、与其他向量存储扩展的关系

本扩展并非 PandasAI 唯一的 VectorStore 实现。在 extensions/ee/vectorstores 目录下，还提供了 ChromaDB（pandasai_chromadb）、Milvus（pandasai_milvus）、Pinecone（pandasai_pinecone）与 Qdrant（pandasai_qdrant）等多个扩展，它们都实现同一套 VectorStore 接口。选择 LanceDB 的典型优势在于：无需部署独立的向量数据库服务，数据以本地文件持久化，通过persist_path即可复用已有训练数据，适合单机与开发环境。

九、许可证与使用边界

最后提醒：pandasai-lancedb采用Sinaptik GmbH Enterprise License（详见 LICENSE 与 README.md），商业场景使用需联系pm@sinaptik.ai获取授权。技术评估与原型验证可参考本文内容，正式商用前请务必确认授权条款。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

企业官网建设流程全解析