将 TiDB 向量检索集成到 Django ORM

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

前置条件

完成本教程，你需要：

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

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

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

运行示例应用

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

步骤 1. 克隆代码仓库

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

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

步骤 2. 创建虚拟环境

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

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

步骤 3. 安装所需依赖

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

pip install -r requirements.txt

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

pip install Django django-tidb mysqlclient numpy python-dotenv

如果在安装 mysqlclient 时遇到问题，请参考 mysqlclient 官方文档。

什么是 django-tidb

django-tidb 是 Django 的 TiDB 方言，增强了 Django ORM 对 TiDB 特性的支持（如向量检索），并解决了 TiDB 与 Django 之间的兼容性问题。

安装 django-tidb 时，请选择与你的 Django 版本相匹配的版本。例如，如果你使用的是 django==4.2.*，则安装 django-tidb==4.2.*。小版本号无需完全一致，建议使用最新的小版本。

更多信息请参考 django-tidb 仓库。

步骤 4. 配置环境变量

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

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

步骤 5. 运行示例

迁移数据库结构：

python manage.py migrate

运行 Django 开发服务器：

python manage.py runserver

在浏览器中访问 http://127.0.0.1:8000 体验示例应用。可用的 API 路径如下：

示例代码片段

你可以参考以下示例代码片段，完成你自己的应用开发。

连接 TiDB 集群

在 sample_project/settings.py 文件中添加如下配置：

dotenv.load_dotenv()

DATABASES = {
    "default": {
        # https://github.com/pingcap/django-tidb
        "ENGINE": "django_tidb",
        "HOST": os.environ.get("TIDB_HOST", "127.0.0.1"),
        "PORT": int(os.environ.get("TIDB_PORT", 4000)),
        "USER": os.environ.get("TIDB_USERNAME", "root"),
        "PASSWORD": os.environ.get("TIDB_PASSWORD", ""),
        "NAME": os.environ.get("TIDB_DATABASE", "test"),
        "OPTIONS": {
            "charset": "utf8mb4",
        },
    }
}

TIDB_CA_PATH = os.environ.get("TIDB_CA_PATH", "")
if TIDB_CA_PATH:
    DATABASES["default"]["OPTIONS"]["ssl_mode"] = "VERIFY_IDENTITY"
    DATABASES["default"]["OPTIONS"]["ssl"] = {
        "ca": TIDB_CA_PATH,
    }

你可以在项目根目录下创建 .env 文件，并将 TIDB_HOST、TIDB_PORT、TIDB_USERNAME、TIDB_PASSWORD、TIDB_DATABASE 和 TIDB_CA_PATH 等环境变量设置为你实际的 TiDB 集群参数。

创建向量表

定义向量列

tidb-django 提供了 VectorField 用于在表中存储向量嵌入。

创建一个包含名为 embedding 的 3 维向量列的表。

class Document(models.Model):
   content = models.TextField()
   embedding = VectorField(dimensions=3)

存储带嵌入向量的文档

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

检索最近邻文档

TiDB 向量检索支持以下距离函数：

• L1Distance
• L2Distance
• CosineDistance
• NegativeInnerProduct

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

results = Document.objects.annotate(
   distance=CosineDistance('embedding', [1, 2, 3])
).order_by('distance')[:3]

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

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

results = Document.objects.annotate(
   distance=CosineDistance('embedding', [1, 2, 3])
).filter(distance__lt=0.2).order_by('distance')[:3]

相关文档

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