TiDB Vector Search を peewee と統合する

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

注記

TiDB Vector Search は現在ベータ版であり、 TiDB Cloudサーバーレスクラスターでのみ使用できます。

前提条件

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

• Python 3.8以上インストールされました。
• ギットインストールされました。
• TiDB Cloud Serverless クラスター。TiDB Cloud クラスターがない場合は、 TiDB Cloud Serverless クラスターの作成に従って独自のTiDB Cloudクラスターを作成してください。

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

以下の手順に従って、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. 環境変数を設定する

1. クラスターページに移動し、ターゲット クラスターの名前をクリックして概要ページに移動します。

2. 右上隅の「接続」をクリックします。接続ダイアログが表示されます。

3. 接続ダイアログの構成が動作環境と一致していることを確認します。

   • 接続タイプはPublicに設定されています。

   • ブランチはmainに設定されています。

   • Connect With はGeneralに設定されています。

   • オペレーティング システムは環境に適合します。

   ヒント：

   プログラムが Windows Subsystem for Linux (WSL) で実行されている場合は、対応する Linux ディストリビューションに切り替えます。

4. 接続ダイアログから接続パラメータをコピーします。

   ヒント：

   まだパスワードを設定していない場合は、 「パスワードの生成」をクリックしてランダムなパスワードを生成します。

5. Python プロジェクトのルート ディレクトリに.envファイルを作成し、接続パラメータを対応する環境変数に貼り付けます。

   • TIDB_HOST : TiDB クラスターのホスト。
   • TIDB_PORT : TiDB クラスターのポート。
   • TIDB_USERNAME : TiDB クラスターに接続するためのユーザー名。
   • TIDB_PASSWORD : TiDB クラスターに接続するためのパスワード。
   • TIDB_DATABASE : 接続するデータベース名。
   • TIDB_CA_PATH : ルート証明書ファイルへのパス。

   以下は macOS の例です。

   TIDB_HOST=gateway01.****.prod.aws.tidbcloud.com
   TIDB_PORT=4000
   TIDB_USERNAME=********.root
   TIDB_PASSWORD=********
   TIDB_DATABASE=test
   TIDB_CA_PATH=/etc/ssl/cert.pem
   

ステップ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)

インデックスで最適化されたベクトル列を定義する

3 次元ベクトル列を定義し、 ベクトル検索インデックス (HNSW インデックス) で最適化します。

class DocumentWithIndex(Model):
    class Meta:
        database = db
        table_name = 'peewee_demo_documents_with_index'

    content = TextField()
    embedding = VectorField(3, constraints=[SQL("COMMENT 'hnsw(distance=cosine)'")])

TiDB はこのインデックスを使用して、コサイン距離関数に基づくベクトル検索クエリを高速化します。

埋め込み付きドキュメントを保存する

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)

参照

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