Skip to content

vector-search-integrate-with-sqlalchemy.md

Latest commit

 

History

History
223 lines (154 loc) · 7.01 KB

vector-search-integrate-with-sqlalchemy.md

File metadata and controls

223 lines (154 loc) · 7.01 KB
title 集成 TiDB 向量搜索与 SQLAlchemy
summary 学习如何将 TiDB 向量搜索与 SQLAlchemy 集成，以存储嵌入向量并执行语义搜索。
aliases
/zh/tidb/stable/vector-search-integrate-with-sqlalchemy/
/zh/tidb/dev/vector-search-integrate-with-sqlalchemy/
/zh/tidbcloud/vector-search-integrate-with-sqlalchemy/

集成 TiDB 向量搜索与 SQLAlchemy

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

注意：

前置条件

完成本教程，你需要：

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

运行示例应用

你可以按照以下步骤快速学习如何将 TiDB 向量搜索与 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 }}} 或 Essential 实例，按以下步骤获取连接字符串并配置环境变量：

  1. 进入 My TiDB 页面，点击目标 {{{ .starter }}} 或 Essential 实例名称进入概览页。

  2. 点击右上角的 Connect，弹出连接对话框。

  3. 确认连接对话框中的配置与你的环境一致。

    • Connection Type 设置为 Public。
    • Branch 设置为 main。
    • Connect With 设置为 SQLAlchemy。
    • Operating System 与你的环境一致。

    提示：

    如果你的程序运行在 Windows Subsystem for Linux (WSL) 中，请切换到对应的 Linux 发行版。

  4. 点击 PyMySQL 标签页并复制连接字符串。

    提示：

    如果你还未设置密码，可点击 Generate Password 生成随机密码。

  5. 在你的 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 集群，在你的 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 {#connect-to-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()

另请参阅