TiDB Vector Search を peewee と統合する
このチュートリアルでは、 ピーウィー使用してTiDB ベクトル検索と対話し、埋め込みを保存し、ベクトル検索クエリを実行する方法について説明します。
注記
TiDB Vector Search は、TiDB Self-Managed (TiDB >= v8.4) およびTiDB Cloudサーバーレスでのみ使用できます。 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. 環境変数を設定する
クラスターページに移動し、ターゲット クラスターの名前をクリックして概要ページに移動します。
右上隅の「接続」をクリックします。接続ダイアログが表示されます。
接続ダイアログの構成が動作環境と一致していることを確認します。
接続タイプは
Publicに設定されています。ブランチは
mainに設定されています。Connect With は
Generalに設定されています。オペレーティング システムは環境に適合します。
ヒント：
プログラムが Windows Subsystem for Linux (WSL) で実行されている場合は、対応する Linux ディストリビューションに切り替えます。
接続ダイアログから接続パラメータをコピーします。
ヒント：
まだパスワードを設定していない場合は、「パスワードの生成」をクリックしてランダムなパスワードを生成します。
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)
埋め込み付きドキュメントを保存する
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)