重要

このページは英語版のページを機械翻訳しています。原文はこちらからご覧ください。

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

このチュートリアルでは、 SQLアルケミー使用して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 を 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. 環境変数を設定する

クラスターページに移動し、ターゲットクラスターの名前をクリックして概要ページに移動します。
右上隅の「接続」をクリックします。接続ダイアログが表示されます。
接続ダイアログの構成が環境と一致していることを確認します。
- 接続タイプはPublicに設定されています。
- ブランチはmainに設定されています。
- Connect With はSQLAlchemyに設定されています。
- オペレーティングシステムは環境に適合します。
ヒント：
プログラムが Windows Subsystem for Linux (WSL) で実行されている場合は、対応する Linux ディストリビューションに切り替えます。
PyMySQLタブをクリックし、接続文字列をコピーします。
ヒント：
まだパスワードを設定していない場合は、「パスワードの生成」をクリックしてランダムなパスワードを生成します。

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"

ステップ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クラスタに接続する

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)

ベクトル列を定義する

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

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()