将 TiDB Vector Search 与 SQLAlchemy 集成
本教程将引导你如何使用 SQLAlchemy 与 TiDB Vector Search 交互,存储嵌入向量,并执行向量搜索查询。
前提条件
完成本教程,你需要:
- 安装 Python 3.8 或更高版本
- 安装 Git
- 一个 TiDB 集群
如果你还没有 TiDB 集群,可以按照以下步骤创建:
- (推荐)参考 创建 TiDB Cloud Starter 集群 来创建你自己的 TiDB Cloud 集群。
- 参考 部署本地测试 TiDB 集群 或 部署生产环境 TiDB 集群 来创建版本为 v8.4.0 或更高版本的本地集群。
运行示例应用
你可以通过以下步骤快速了解如何将 TiDB Vector Search 与 SQLAlchemy 集成。
第 1 步:克隆仓库
将 tidb-vector-python
仓库克隆到你的本地机器:
git clone https://github.com/pingcap/tidb-vector-python.git
第 2 步:创建虚拟环境
为你的项目创建一个虚拟环境:
cd tidb-vector-python/examples/orm-sqlalchemy-quickstart
python3 -m venv .venv
source .venv/bin/activate
第 3 步:安装所需依赖
安装示例项目所需的依赖:
pip install -r requirements.txt
或者,你也可以单独安装以下包:
pip install pymysql python-dotenv sqlalchemy tidb-vector
第 4 步:配置环境变量
根据你选择的 TiDB 部署方式,配置环境变量。
对于 TiDB Cloud Starter 集群,按照以下步骤获取集群连接字符串并配置环境变量:
进入 Clusters 页面,然后点击目标集群的名称,进入其概览页面。
点击右上角的 Connect,弹出连接对话框。
确认连接对话框中的配置与环境一致。
- Connection Type 设置为
Public
。 - Branch 设置为
main
。 - Connect With 设置为
SQLAlchemy
。 - Operating System 与你的环境匹配。
- Connection Type 设置为
点击 PyMySQL 标签页,复制连接字符串。
在你的 Python 项目的根目录下,创建
.env
文件,并将连接字符串粘贴进去。下面是 macOS 的示例:
TIDB_DATABASE_URL="mysql+pymysql://<prefix>.root:<password>@gateway01.<region>.prod.aws.tidbcloud.com:4000/test?ssl_ca=/etc/ssl/cert.pem&ssl_verify_cert=true&ssl_verify_identity=true"
对于 TiDB Self-Managed 集群,在你的 Python 项目的根目录下创建 .env
文件。将以下内容复制到 .env
文件中,并根据你的 TiDB 集群连接参数修改环境变量值:
TIDB_DATABASE_URL="mysql+pymysql://<USER>:<PASSWORD>@<HOST>:<PORT>/<DATABASE>"
# 例如:TIDB_DATABASE_URL="mysql+pymysql://root@127.0.0.1:4000/test"
如果你在本地运行 TiDB,<HOST>
默认为 127.0.0.1
。初次启动集群时,<PASSWORD>
默认为空,可以省略。
以下是各参数的说明:
<USER>
:连接 TiDB 集群的用户名。<PASSWORD>
:连接 TiDB 集群的密码。<HOST>
:TiDB 集群的主机地址。<PORT>
:TiDB 集群的端口。<DATABASE>
:你要连接的数据库名称。
第 5 步:运行示例
python sqlalchemy-quickstart.py
示例输出:
Get 3-nearest neighbor documents:
- distance: 0.00853986601633272
document: fish
- distance: 0.12712843905603044
document: dog
- distance: 0.7327387580875756
document: tree
Get documents within a certain distance:
- distance: 0.00853986601633272
document: fish
- distance: 0.12712843905603044
document: dog
示例代码片段
你可以参考以下示例代码片段,开发你的应用。
创建向量表
连接到 TiDB 集群
import os
import dotenv
from sqlalchemy import Column, Integer, create_engine, Text
from sqlalchemy.orm import declarative_base, Session
from tidb_vector.sqlalchemy import VectorType
dotenv.load_dotenv()
tidb_connection_string = os.environ['TIDB_DATABASE_URL']
engine = create_engine(tidb_connection_string)
定义向量列
创建一个表,包含名为 embedding
的列,用于存储 3 维向量。
Base = declarative_base()
class Document(Base):
__tablename__ = 'sqlalchemy_demo_documents'
id = Column(Integer, primary_key=True)
content = Column(Text)
embedding = Column(VectorType(3))
存储带有嵌入向量的文档
with Session(engine) as session:
session.add(Document(content="dog", embedding=[1, 2, 1]))
session.add(Document(content="fish", embedding=[1, 2, 4]))
session.add(Document(content="tree", embedding=[1, 0, 0]))
session.commit()
搜索最近邻文档
根据余弦距离函数,搜索与查询向量 [1, 2, 3]
语义最接近的前 3 个文档。
with Session(engine) as session:
distance = Document.embedding.cosine_distance([1, 2, 3]).label('distance')
results = session.query(
Document, distance
).order_by(distance).limit(3).all()
搜索距离在一定范围内的文档
搜索与查询向量 [1, 2, 3]
的余弦距离小于 0.2 的文档。
with Session(engine) as session:
distance = Document.embedding.cosine_distance([1, 2, 3]).label('distance')
results = session.query(
Document, distance
).filter(distance < 0.2).order_by(distance).limit(3).all()