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

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

注記：

• ベクター検索機能はベータ版です。予告なく変更される場合があります。バグを発見した場合は、GitHubで問題を報告してください。
• ベクトル検索機能は、 TiDBセルフマネージドTiDB Cloud Starter 、 TiDB Cloud Essential 、およびTiDB Cloud Dedicatedで利用できます。TiDB Self-ManagedおよびTiDB Cloud Dedicatedの場合、TiDBのバージョンはv8.4.0以降である必要があります（v8.5.0以降を推奨）。

前提条件

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

• Python 3.8以降インストールされています。
• Gitがインストールされました。
• TiDBクラスタ。

TiDBクラスタをお持ちでない場合は、以下の手順で作成できます。

• (推奨) TiDB Cloud Starterインスタンスを作成する。
• ローカルテスト用のTiDBセルフマネージドクラスタをデプロイ。または本番本番のTiDBセルフマネージドクラスタをデプロイ。

サンプルアプリを実行します

以下の手順に従うことで、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,
)

ベクトル列を定義する

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

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)

関連項目

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