Python を使って TiDB + AI を始めよう

このドキュメントでは、Python SDK を使用して TiDB でベクトル検索開始する方法を説明します。手順に従って、TiDB で動作する最初の AI アプリケーションを構築します。

このドキュメントに従うことで、以下のことを学ぶことができます。

• TiDB Python SDKを使用してTiDBに接続します。
• 一般的な埋め込みモデルを使用してテキスト埋め込みを生成します。
• ベクトルを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以降を推奨）。

前提条件

• tidbcloud.comにアクセスしてTiDB Cloud Starterインスタンスを無料で作成するか、 ティアップ遊び場を使用してローカル テスト用の TiDB 自己管理クラスターをデプロイします。

インストール

pytidbはTiDBの公式Python SDKであり、開発者がAIアプリケーションを効率的に構築できるよう設計されています。

Python SDKをインストールするには、次のコマンドを実行してください。

pip install pytidb

組み込みの埋め込み機能を使用するには、 models拡張機能をインストールしてください（代替案）：

pip install "pytidb[models]"

データベースに接続します

HOST:     gateway01.us-east-1.prod.shared.aws.tidbcloud.com
PORT:     4000
USERNAME: 4EfqPF23YKBxaQb.root
PASSWORD: abcd1234
DATABASE: test
CA:       /etc/ssl/cert.pem

from pytidb import TiDBClient

client = TiDBClient.connect(
    host="gateway01.us-east-1.prod.shared.aws.tidbcloud.com",
    port=4000,
    username="4EfqPF23YKBxaQb.root",
    password="abcd1234",
    database="test",
)

from pytidb import TiDBClient

client = TiDBClient.connect(
    host="localhost",
    port=4000,
    username="root",
    password="",
    database="test",
    ensure_db=True,
)

接続が完了すると、 clientオブジェクトを使用して、テーブルの操作、データのクエリなどを行うことができます。

埋め込み関数を作成する

埋め込みモデル扱う場合、埋め込み機能を利用することで、挿入時とクエリ時の両方でデータを自動的にベクトル化できます。OpenAI、Jina AI、Hugging Face、Sentence Transformersなど、人気の埋め込みモデルをネイティブでサポートしています。

from pytidb.embeddings import EmbeddingFunction

text_embed = EmbeddingFunction(
    model_name="openai/text-embedding-3-small",
    api_key="<your-openai-api-key>",
)

from pytidb.embeddings import EmbeddingFunction

text_embed = EmbeddingFunction(
    model_name="jina/jina-embeddings-v3",
    api_key="<your-jina-api-key>",
)

テーブルを作成する

例として、 chunksという名前のテーブルを作成し、以下の列を追加します。

• id (int): チャンクのID。
• text (テキスト): チャンクのテキストコンテンツ。
• text_vec (ベクトル): テキストのベクトル埋め込み。
• user_id (int): チャンクを作成したユーザーのID。

from pytidb.schema import TableModel, Field, VectorField

class Chunk(TableModel):
    id: int | None = Field(default=None, primary_key=True)
    text: str = Field()
    text_vec: list[float] = text_embed.VectorField(source_field="text")
    user_id: int = Field()

table = client.create_table(schema=Chunk, if_exists="overwrite")

作成後は、 tableオブジェクトを使用して、データの挿入、データの検索などを行うことができます。

データを挿入する

それでは、サンプルデータをテーブルに追加してみましょう。

table.bulk_insert([
    # 👇 The text will be automatically embedded and populated into the `text_vec` field.
    Chunk(text="PyTiDB is a Python library for developers to connect to TiDB.", user_id=2),
    Chunk(text="LlamaIndex is a framework for building AI applications.", user_id=2),
    Chunk(text="OpenAI is a company and platform that provides AI models service and tools.", user_id=3),
])

近隣住民を検索

指定されたクエリの最近傍を検索するには、 table.search()メソッドを使用できます。このメソッドは、 デフォルトでベクトル検索実行します。

table.search(
    # 👇 Pass the query text directly, it will be embedded to a query vector automatically.
    "A library for my artificial intelligence software"
)
.limit(3).to_list()

この例では、ベクトル検索はクエリベクトルtext_vecテーブルのchunksフィールドに格納されているベクトルと比較し、類似度スコアに基づいて意味的に最も関連性の高い上位3つの結果を返します。

_distanceが近いほど、2 つのベクトルはより類似していることを意味します。

[
    {
        'id': 2,
        'text': 'LlamaIndex is a framework for building AI applications.',
        'text_vec': [...],
        'user_id': 2,
        '_distance': 0.5719928358786761,
        '_score': 0.4280071641213239
    },
    {
        'id': 3,
        'text': 'OpenAI is a company and platform that provides AI models service and tools.',
        'text_vec': [...],
        'user_id': 3,
        '_distance': 0.603133726213383,
        '_score': 0.396866273786617
    },
    {
        'id': 1,
        'text': 'PyTiDB is a Python library for developers to connect to TiDB.',
        'text_vec': [...],
        'user_id': 2,
        '_distance': 0.6202191842385758,
        '_score': 0.3797808157614242
    }
]

データを削除する

テーブルから特定の行を削除するには、 table.delete()メソッドを使用できます。

table.delete({
    "id": 1
})

ドロップテーブル

テーブルが不要になった場合は、 client.drop_table()メソッドを使用して削除できます。

client.drop_table("chunks")

次のステップ

• TiDB のベクトル検索、 全文検索、 ハイブリッド検索の詳細については、こちらをご覧ください。