TiDBベクトル検索をJina AI埋め込みAPIと統合する

このチュートリアルではジナAIを使用してテキスト埋め込みを生成し、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をJina AIの埋め込み機能と統合する方法をすぐに習得できます。

ステップ1. リポジトリをクローンする

tidb-vector-pythonリポジトリをローカルマシンにクローンします。

git clone https://github.com/pingcap/tidb-vector-python.git

ステップ2. 仮想環境を作成する

プロジェクト用の仮想環境を作成する：

cd tidb-vector-python/examples/jina-ai-embeddings-demo
python3 -m venv .venv
source .venv/bin/activate

ステップ3. 必要な依存関係をインストールします

デモプロジェクトに必要な依存関係をインストールします。

pip install -r requirements.txt

ステップ4．環境変数を設定する

Jina AIのAPIキーをJina AI 埋め込み APIページから取得し、選択したTiDBデプロイメントオプションに応じて環境変数を設定してください。

export JINA_API_KEY="****"
export TIDB_DATABASE_URL="mysql+pymysql://<USERNAME>:<PASSWORD>@<HOST>:<PORT>/<DATABASE>"
# For example: export TIDB_DATABASE_URL="mysql+pymysql://root@127.0.0.1:4000/test"

ステップ5．デモを実行する

python jina-ai-embeddings-demo.py

出力例：

- Inserting Data to TiDB...
  - Inserting: Jina AI offers best-in-class embeddings, reranker and prompt optimizer, enabling advanced multimodal AI.
  - Inserting: TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.
- List All Documents and Their Distances to the Query:
  - distance: 0.3585317326132522
    content: Jina AI offers best-in-class embeddings, reranker and prompt optimizer, enabling advanced multimodal AI.
  - distance: 0.10858102967720984
    content: TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.
- The Most Relevant Document and Its Distance to the Query:
  - distance: 0.10858102967720984
    content: TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.

サンプルコードスニペット

Jina AIから埋め込みデータを取得する

generate_embeddings AI埋め込みAPIを呼び出すためのヘルパー関数を定義します。

import os
import requests
import dotenv

dotenv.load_dotenv()

JINAAI_API_KEY = os.getenv('JINAAI_API_KEY')

def generate_embeddings(text: str):
    JINAAI_API_URL = 'https://api.jina.ai/v1/embeddings'
    JINAAI_HEADERS = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {JINAAI_API_KEY}'
    }
    JINAAI_REQUEST_DATA = {
        'input': [text],
        'model': 'jina-embeddings-v2-base-en'  # with dimension 768.
    }
    response = requests.post(JINAAI_API_URL, headers=JINAAI_HEADERS, json=JINAAI_REQUEST_DATA)
    return response.json()['data'][0]['embedding']

TiDBに接続する

SQLAlchemy経由でTiDBに接続する：

import os
import dotenv

from tidb_vector.sqlalchemy import VectorType
from sqlalchemy.orm import Session, declarative_base

dotenv.load_dotenv()

TIDB_DATABASE_URL = os.getenv('TIDB_DATABASE_URL')
assert TIDB_DATABASE_URL is not None
engine = create_engine(url=TIDB_DATABASE_URL, pool_recycle=300)

ベクトルテーブルのスキーマを定義します

jinaai_tidb_demo_documentsという名前のテーブルを作成し、テキストを格納するためのcontent列と、埋め込みを格納するためのcontent_vecという名前のベクトル列を作成します。

from sqlalchemy import Column, Integer, String, create_engine
from sqlalchemy.orm import declarative_base

Base = declarative_base()

class Document(Base):
    __tablename__ = "jinaai_tidb_demo_documents"

    id = Column(Integer, primary_key=True)
    content = Column(String(255), nullable=False)
    content_vec = Column(
        # DIMENSIONS is determined by the embedding model,
        # for Jina AI's jina-embeddings-v2-base-en model it's 768.
        VectorType(dim=768),
        comment="hnsw(distance=cosine)"

注記：

• ベクトル列の次元は、埋め込みモデルによって生成される埋め込みの次元と一致していなければなりません。
• この例では、 jina-embeddings-v2-base-enモデルによって生成される埋め込みの次元は768です。

Jina AIで埋め込みを作成し、TiDBに保存します。

Jina AI Embeddings APIを使用して、各テキストの埋め込みを生成し、その埋め込みをTiDBに保存します。

TEXTS = [
   'Jina AI offers best-in-class embeddings, reranker and prompt optimizer, enabling advanced multimodal AI.',
   'TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.',
]
data = []

for text in TEXTS:
    # Generate embeddings for the texts via Jina AI API.
    embedding = generate_embeddings(text)
    data.append({
        'text': text,
        'embedding': embedding
    })

with Session(engine) as session:
   print('- Inserting Data to TiDB...')
   for item in data:
      print(f'  - Inserting: {item["text"]}')
      session.add(Document(
         content=item['text'],
         content_vec=item['embedding']
      ))
   session.commit()

TiDBでJina AI埋め込みを用いたセマンティック検索を実行する

Jina AIの埋め込みAPIを使用してクエリテキストの埋め込みを生成し、クエリテキストの埋め込みとベクトルテーブル内の各埋め込みとの間のコサイン距離に基づいて最も関連性の高いドキュメントを検索します。

query = 'What is TiDB?'
# Generate the embedding for the query via Jina AI API.
query_embedding = generate_embeddings(query)

with Session(engine) as session:
    print('- The Most Relevant Document and Its Distance to the Query:')
    doc, distance = session.query(
        Document,
        Document.content_vec.cosine_distance(query_embedding).label('distance')
    ).order_by(
        'distance'
    ).limit(1).first()
    print(f'  - distance: {distance}\n'
          f'    content: {doc.content}')

関連項目

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