TiDBベクトル検索をpeeweeと統合する

このチュートリアルでは、 ピーウィー使用してTiDBベクトル検索と対話し、埋め込みを保存し、ベクトル検索クエリを実行する方法について説明します。

注記：

ベクター検索機能はベータ版です。予告なく変更される可能性があります。バグを見つけた場合は、GitHubで問題報告を行ってください。

注記：

ベクトル検索機能は、TiDB Self-Managed、 TiDB Cloudスターター TiDB Cloudエッセンシャル利用できますTiDB Cloud専用 Self-Managed およびTiDB Cloud Dedicated の場合、TiDB バージョンは v8.4.0 以降である必要があります（v8.5.0 以降を推奨）。

前提条件

このチュートリアルを完了するには、次のものが必要です。

• Python 3.8以上個インストールされました。
• ギット個インストールされました。
• TiDB クラスター。

TiDB クラスターがない場合は、次のように作成できます。

• (推奨) TiDB Cloud Starter クラスターの作成に従って、独自のTiDB Cloudクラスターを作成します。
• ローカルテストTiDBクラスタをデプロイまたは本番のTiDBクラスタをデプロイに従って、v8.4.0 以降のバージョンのローカル クラスターを作成します。

サンプルアプリを実行する

以下の手順に従って、TiDB Vector Search を 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,
)

ベクトル列を定義する

3 次元ベクトルを格納するpeewee_demo_documentsという名前の列を持つテーブルを作成します。

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)

参照

• ベクトルデータ型
• ベクター検索インデックス