TiDB Vector Search を Amazon Bedrock に統合する
このチュートリアルでは、TiDB のベクトル検索機能をアマゾンの岩盤と統合して、検索拡張生成 (RAG) Q&A ボットを構築する方法を説明します。
注記：
ベクター検索機能はベータ版です。予告なく変更される可能性があります。バグを発見した場合は、GitHubで問題報告を行ってください。
注記
ベクトル検索機能は、TiDB Self-Managed、 TiDB Cloudスターター TiDB Cloudエッセンシャル利用できますTiDB Cloud専用 Self-ManagedおよびTiDB Cloud Dedicatedの場合、TiDBバージョンはv8.4.0以降である必要があります（v8.5.0以降を推奨）。
ヒント
サンプルコード全体をノートブック形式で閲覧できます。
前提条件
このチュートリアルを完了するには、次のものが必要です。
Python 3.11以降インストール済み
ピップインストール済み
AWS CLIインストール済み
このチュートリアルでは、AWS CLI プロファイルがサポートされているアマゾンの岩盤リージョンに設定されていることを確認してください。サポートされているリージョンの一覧はアマゾンの岩盤地域にあります。サポートされているリージョンに切り替えるには、次のコマンドを実行します。
aws configure set region <your-region>TiDB Cloud Starter クラスター
独自のTiDB Cloudクラスターがない場合は、手順TiDB Cloud Starter クラスターの作成に従って作成します。
Amazon Bedrock に必要な権限と次のモデルへのアクセス権を持つ AWS アカウント:
- Amazon Titan Embeddings (
amazon.titan-embed-text-v2:0) テキスト埋め込みの生成に使用 - Meta Llama 3 (
us.meta.llama3-2-3b-instruct-v1:0)、テキスト生成に使用
アクセスできない場合は、 Amazon Bedrock 基盤モデルへのアクセスをリクエストするの手順に従ってください。
- Amazon Titan Embeddings (
始めましょう
このセクションでは、TiDB Vector Search を Amazon Bedrock と統合して RAG ベースの Q&A ボットを構築するための手順を段階的に説明します。
ステップ1.環境変数を設定する
TiDB Cloudコンソールから TiDB 接続情報を取得し、開発環境で環境変数を次のように設定します。
クラスターページに移動し、ターゲット クラスターの名前をクリックして概要ページに移動します。
右上隅の「接続」をクリックします。接続ダイアログが表示されます。
接続ダイアログの構成が動作環境と一致していることを確認します。
接続タイプは
Publicに設定されていますブランチは
mainに設定されています接続先が
Generalに設定されていますオペレーティング システムは環境に適合します。
ヒント：
プログラムが Windows Subsystem for Linux (WSL) で実行されている場合は、対応する Linux ディストリビューションに切り替えます。
ランダムなパスワードを作成するには、 「パスワードの生成」をクリックします。
ヒント：
以前にパスワードを作成したことがある場合は、元のパスワードを使用するか、 「パスワードのリセット」をクリックして新しいパスワードを生成することができます。
環境変数を設定するには、ターミナルで以下のコマンドを実行します。コマンド内のプレースホルダーは、接続ダイアログから取得した対応する接続パラメータに置き換える必要があります。
export TIDB_HOST=<your-tidb-host> export TIDB_PORT=4000 export TIDB_USER=<your-tidb-user> export TIDB_PASSWORD=<your-tidb-password> export TIDB_DB_NAME=test
ステップ2. Python仮想環境を設定する
demo.pyという名前の Python ファイルを作成します。touch demo.py依存関係を管理するために仮想環境を作成してアクティブ化します。
python3 -m venv env source env/bin/activate # On Windows, use env\Scripts\activate必要な依存関係をインストールします。
pip install SQLAlchemy==2.0.30 PyMySQL==1.1.0 tidb-vector==0.0.9 pydantic==2.7.1 boto3
ステップ3. 必要なライブラリをインポートする
必要なライブラリをインポートするには、 demo.pyの先頭に次のコードを追加します。
import os
import json
import boto3
from sqlalchemy import Column, Integer, Text, create_engine
from sqlalchemy.orm import declarative_base, Session
from tidb_vector.sqlalchemy import VectorType
ステップ4. データベース接続を構成する
demo.pyで、次のコードを追加してデータベース接続を構成します。
# ---- Configuration Setup ----
# Set environment variables: TIDB_HOST, TIDB_PORT, TIDB_USER, TIDB_PASSWORD, TIDB_DB_NAME
TIDB_HOST = os.environ.get("TIDB_HOST")
TIDB_PORT = os.environ.get("TIDB_PORT")
TIDB_USER = os.environ.get("TIDB_USER")
TIDB_PASSWORD = os.environ.get("TIDB_PASSWORD")
TIDB_DB_NAME = os.environ.get("TIDB_DB_NAME")
# ---- Database Setup ----
def get_db_url():
"""Build the database connection URL."""
return f"mysql+pymysql://{TIDB_USER}:{TIDB_PASSWORD}@{TIDB_HOST}:{TIDB_PORT}/{TIDB_DB_NAME}?ssl_verify_cert=True&ssl_verify_identity=True"
# Create engine
engine = create_engine(get_db_url(), pool_recycle=300)
Base = declarative_base()
ステップ 5. Bedrock ランタイム クライアントを使用して Amazon Titan Text Embeddings V2 モデルを呼び出す
Amazon Bedrock ランタイムクライアントは、次のパラメータを受け入れるinvoke_model API を提供します。
modelId: Amazon Bedrock で利用可能な基盤モデルのモデル ID。accept: 入力要求のタイプ。contentType: 入力のコンテンツ タイプ。body: プロンプトと構成で構成される JSON 文字列ペイロード。
demo.pyに次のコードを追加してinvoke_model API を呼び出し、Amazon Titan Text Embeddings を使用してテキスト埋め込みを生成し、Meta Llama 3 からの応答を取得します。
# Bedrock Runtime Client Setup
bedrock_runtime = boto3.client('bedrock-runtime')
# ---- Model Invocation ----
embedding_model_name = "amazon.titan-embed-text-v2:0"
dim_of_embedding_model = 512
llm_name = "us.meta.llama3-2-3b-instruct-v1:0"
def embedding(content):
"""Invoke Amazon Bedrock to get text embeddings."""
payload = {
"modelId": embedding_model_name,
"contentType": "application/json",
"accept": "*/*",
"body": {
"inputText": content,
"dimensions": dim_of_embedding_model,
"normalize": True,
}
}
body_bytes = json.dumps(payload['body']).encode('utf-8')
response = bedrock_runtime.invoke_model(
body=body_bytes,
contentType=payload['contentType'],
accept=payload['accept'],
modelId=payload['modelId']
)
result_body = json.loads(response.get("body").read())
return result_body.get("embedding")
def generate_result(query: str, info_str: str):
"""Generate answer using Meta Llama 3 model."""
prompt = f"""
ONLY use the content below to generate an answer:
{info_str}
----
Please carefully think about the question: {query}
"""
payload = {
"modelId": llm_name,
"contentType": "application/json",
"accept": "application/json",
"body": {
"prompt": prompt,
"temperature": 0
}
}
body_bytes = json.dumps(payload['body']).encode('utf-8')
response = bedrock_runtime.invoke_model(
body=body_bytes,
contentType=payload['contentType'],
accept=payload['accept'],
modelId=payload['modelId']
)
result_body = json.loads(response.get("body").read())
completion = result_body["generation"]
return completion
ステップ6.ベクターテーブルを作成する
demo.pyに次のコードを追加して、テキストとベクトル埋め込みを格納するベクトル テーブルを作成します。
# ---- TiDB Setup and Vector Index Creation ----
class Entity(Base):
"""Define the Entity table with a vector index."""
__tablename__ = "entity"
id = Column(Integer, primary_key=True)
content = Column(Text)
content_vec = Column(VectorType(dim=dim_of_embedding_model), comment="hnsw(distance=l2)")
# Create the table in TiDB
Base.metadata.create_all(engine)
ステップ7. ベクターデータをTiDB Cloud Starterに保存する
demo.pyで、次のコードを追加して、ベクター データをTiDB Cloud Starter クラスターに保存します。
# ---- Saving Vectors to TiDB ----
def save_entities_with_embedding(session, contents):
"""Save multiple entities with their embeddings to the TiDB database."""
for content in contents:
entity = Entity(content=content, content_vec=embedding(content))
session.add(entity)
session.commit()
ステップ8. アプリケーションを実行する
demo.pyでは、次のコードを追加して、データベース セッションを確立し、埋め込みを TiDB に保存し、サンプルの質問 (「TiDB とは何ですか?」など) を尋ね、モデルから結果を生成します。if __name__ == "__main__": # Establish a database session with Session(engine) as session: # Example data contents = [ "TiDB is a distributed SQL database compatible with MySQL.", "TiDB supports Hybrid Transactional and Analytical Processing (HTAP).", "TiDB can scale horizontally and provides high availability.", "Amazon Bedrock allows seamless integration with foundation models.", "Meta Llama 3 is a powerful model for text generation." ] # Save embeddings to TiDB save_entities_with_embedding(session, contents) # Example query query = "What is TiDB?" info_str = " ".join(contents) # Generate result from Meta Llama 3 result = generate_result(query, info_str) print(f"Generated answer: {result}")すべての変更を
demo.pyに保存し、スクリプトを実行します。python3 demo.py予想される出力は次のようになります。
Generated answer: What is the main purpose of TiDB? What are the key features of TiDB? What are the key benefits of TiDB? ---- Based on the provided text, here is the answer to the question: What is TiDB? TiDB is a distributed SQL database compatible with MySQL. ## Step 1: Understand the question The question asks for the definition of TiDB. ## Step 2: Identify the key information The key information provided in the text is that TiDB is a distributed SQL database compatible with MySQL. ## Step 3: Provide the answer Based on the provided text, TiDB is a distributed SQL database compatible with MySQL. The final answer is: TiDB is a distributed SQL database compatible with MySQL.