1. 案例目标
本案例展示了如何使用TiDB Cloud的向量搜索功能与LlamaIndex集成,实现高效的文档检索和语义搜索。TiDB Serverless将内置的向量搜索集成到MySQL生态系统中,使用户无需额外的数据库或技术栈即可开发AI应用程序。
2. 技术栈与核心依赖
LlamaIndex
TiDB Cloud
OpenAI API
llama-index-vector-stores-tidbvector
PyMySQL
SimpleDirectoryReader
VectorStoreIndex
3. 环境配置
安装必要依赖
%pip install llama-index-vector-stores-tidbvector
%pip install llama-index
配置OpenAI API密钥
import getpass
import os
os.environ["OPENAI_API_KEY"] = getpass.getpass("Input your OpenAI API key:")
配置TiDB连接字符串
# 从TiDB Cloud控制台获取连接字符串
tidb_connection_string_template = "mysql+pymysql://<USER>:<PASSWORD>@<HOST>:4000/<DB>?ssl_ca=/etc/ssl/cert.pem&ssl_verify_cert=true&ssl_verify_identity=true"
tidb_password = getpass.getpass("Input your TiDB password:")
tidb_connection_url = tidb_connection_string_template.replace("<PASSWORD>", tidb_password)
4. 案例实现
4.1 数据准备
下载并加载Paul Graham文章
!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'
documents = SimpleDirectoryReader("./data/paul_graham").load_data()
for index, document in enumerate(documents):
document.metadata = {"book": "paul_graham"}
4.2 创建TiDB向量存储
初始化TiDBVectorStore
VECTOR_TABLE_NAME = "paul_graham_test"
tidbvec = TiDBVectorStore(
connection_string=tidb_connection_url,
table_name=VECTOR_TABLE_NAME,
distance_strategy="cosine",
vector_dimension=1536,
drop_existing_table=False,
)
4.3 构建向量索引
使用TiDB向量存储创建索引
storage_context = StorageContext.from_defaults(vector_store=tidbvec)
index = VectorStoreIndex.from_documents(
documents, storage_context=storage_context, show_progress=True
)
注意:如果由于MySQL协议的数据包大小限制在插入大量向量(如2000行)时遇到错误,可以通过设置更小的insert_batch_size参数(例如1000)来缓解此问题,确保数据顺利插入TiDB向量存储。
4.4 语义相似性搜索
基本查询
query_engine = index.as_query_engine()
response = query_engine.query("What did the author do?")
print(textwrap.fill(str(response), 100))
4.5 元数据过滤
使用元数据过滤器进行搜索
from llama_index.core.vector_stores.types import (
MetadataFilter,
MetadataFilters,
)
# 使用不等于过滤器
query_engine = index.as_query_engine(
filters=MetadataFilters(
filters=[
MetadataFilter(key="book", value="paul_graham", operator="!="),
]
),
similarity_top_k=2,
)
response = query_engine.query("What did the author learn?")
# 使用等于过滤器
query_engine = index.as_query_engine(
filters=MetadataFilters(
filters=[
MetadataFilter(key="book", value="paul_graham", operator="=="),
]
),
similarity_top_k=2,
)
response = query_engine.query("What did the author learn?")
4.6 文档删除
从向量存储中删除文档
tidbvec.delete(documents[0].doc_id)
5. 案例效果
- 成功将Paul Graham的文章加载到TiDB向量存储中
- 实现了基于语义相似性的文档检索
- 通过元数据过滤器精确控制搜索结果
- 支持文档的删除操作
- 展示了TiDB向量存储与LlamaIndex的无缝集成
6. 案例实现思路
- 环境准备:安装必要的Python包,配置API密钥和数据库连接
- 数据加载:使用SimpleDirectoryReader加载文档数据,并为文档添加元数据
- 向量存储初始化:创建TiDBVectorStore实例,配置表名、距离策略和向量维度
- 索引构建:使用TiDB向量存储创建VectorStoreIndex,将文档向量化存储
- 查询处理:创建查询引擎,执行基本查询和带过滤条件的查询
- 文档管理:演示如何从向量存储中删除文档
7. 扩展建议
- 批量操作优化:对于大规模数据,可以进一步优化批量插入和查询性能
- 高级过滤:探索更复杂的元数据过滤条件,如范围查询、多条件组合
- 混合搜索:结合向量搜索和传统关键词搜索,提高搜索准确性
- 性能监控:添加查询性能监控,优化索引配置
- 多模态支持:扩展支持图像、音频等多模态数据的向量存储和检索
- 实时更新:实现向量存储的实时更新机制,保持数据同步
8. 总结
TiDB向量存储示例展示了如何将传统关系型数据库与现代向量搜索技术结合,为AI应用提供强大的数据检索能力。通过TiDB Cloud的向量搜索功能,开发者可以在熟悉的MySQL生态系统中实现高效的语义搜索,无需额外的技术栈。这种集成方案特别适合已有MySQL技术基础、希望快速构建AI应用的企业和开发者。