将 TiDB 向量检索集成到 peewee

本教程将指导你如何使用 peewee 与 TiDB 向量检索 进行交互，存储嵌入向量，并执行向量检索查询。

前置条件

完成本教程，你需要：

• 已安装 Python 3.8 或更高版本。
• 已安装 Git。
• 一个 TiDB 集群。

如果你还没有 TiDB 集群，可以按如下方式创建：

• （推荐）参考 创建 TiDB Cloud Starter 集群 创建属于你自己的 TiDB Cloud 集群。
• 参考 部署本地测试 TiDB 集群 或 部署生产环境 TiDB 集群 创建 v8.4.0 或更高版本的本地集群。

运行示例应用

你可以按照以下步骤快速了解如何将 TiDB 向量检索与 peewee 集成。

步骤 1. 克隆代码仓库

将 tidb-vector-python 仓库克隆到本地：

git clone https://github.com/pingcap/tidb-vector-python.git

步骤 2. 创建虚拟环境

为你的项目创建一个虚拟环境：

cd tidb-vector-python/examples/orm-peewee-quickstart
python3 -m venv .venv
source .venv/bin/activate

步骤 3. 安装所需依赖

为示例项目安装所需依赖：

pip install -r requirements.txt

或者，你也可以为你的项目单独安装以下包：

pip install peewee pymysql python-dotenv tidb-vector

步骤 4. 配置环境变量

根据你选择的 TiDB 部署方式配置环境变量。

TIDB_HOST=127.0.0.1
TIDB_PORT=4000
TIDB_USERNAME=root
TIDB_PASSWORD=
TIDB_DATABASE=test

步骤 5. 运行示例

python peewee-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 peewee import Model, MySQLDatabase, SQL, TextField
from tidb_vector.peewee import VectorField

dotenv.load_dotenv()

# Using `pymysql` as the driver.
connect_kwargs = {
    'ssl_verify_cert': True,
    'ssl_verify_identity': True,
}

# Using `mysqlclient` as the driver.
# connect_kwargs = {
#     'ssl_mode': 'VERIFY_IDENTITY',
#     'ssl': {
#         # Root certificate default path
#         # https://docs.pingcap.com/tidbcloud/secure-connections-to-serverless-clusters/#root-certificate-default-path
#         'ca': os.environ.get('TIDB_CA_PATH', '/path/to/ca.pem'),
#     },
# }

db = MySQLDatabase(
    database=os.environ.get('TIDB_DATABASE', 'test'),
    user=os.environ.get('TIDB_USERNAME', 'root'),
    password=os.environ.get('TIDB_PASSWORD', ''),
    host=os.environ.get('TIDB_HOST', 'localhost'),
    port=int(os.environ.get('TIDB_PORT', '4000')),
    **connect_kwargs,
)

定义向量列

创建一个包含名为 peewee_demo_documents 的表，并存储 3 维向量。

class Document(Model):
    class Meta:
        database = db
        table_name = 'peewee_demo_documents'

    content = TextField()
    embedding = VectorField(3)

存储带嵌入向量的文档

Document.create(content='dog', embedding=[1, 2, 1])
Document.create(content='fish', embedding=[1, 2, 4])
Document.create(content='tree', embedding=[1, 0, 0])

检索最近邻文档

基于余弦距离函数，检索与查询向量 [1, 2, 3] 语义最接近的前 3 个文档。

distance = Document.embedding.cosine_distance([1, 2, 3]).alias('distance')
results = Document.select(Document, distance).order_by(distance).limit(3)

检索距离在指定范围内的文档

检索与查询向量 [1, 2, 3] 的余弦距离小于 0.2 的文档。

distance_expression = Document.embedding.cosine_distance([1, 2, 3])
distance = distance_expression.alias('distance')
results = Document.select(Document, distance).where(distance_expression < 0.2).order_by(distance).limit(3)

相关阅读

• 向量数据类型
• 向量检索索引