重要

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

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

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

注記
TiDB Vector Search は現在ベータ版であり、 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))

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

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

class DocumentWithIndex(Base):
    __tablename__ = 'sqlalchemy_demo_documents_with_index'
    id = Column(Integer, primary_key=True)
    content = Column(Text)
    embedding = Column(VectorType(3), comment="hnsw(distance=cosine)")

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

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

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