TiDB Vector SearchをAmazon Bedrockと統合する
注記：
このドキュメントはTiDB Cloudにのみ適用され、TiDB Self-Managedには適用されません。
このチュートリアルではTiDB ベクトル検索とアマゾンの岩盤統合して、検索拡張生成 (RAG) Q&A ボットを構築する方法を説明します。
注記：
- ベクター検索機能はベータ版です。予告なく変更される場合があります。バグを発見した場合は、GitHubで問題を報告してください。
- ベクトル検索機能は、 TiDBセルフマネージドTiDB Cloud Starter 、 TiDB Cloud Essential 、およびTiDB Cloud Dedicatedで利用できます。TiDB Self-ManagedおよびTiDB Cloud Dedicatedの場合、TiDBのバージョンはv8.4.0以降である必要があります（v8.5.0以降を推奨）。
ヒント
完全なサンプルコードNotebook 形式で表示できます。
前提条件
このチュートリアルを完了するには、以下が必要です。
Python 3.11以降インストールされています
ピップがインストールされました
AWS CLIがインストールされました
AWS CLI プロファイルがサポートされているアマゾンの岩盤リージョンに設定されていることを確認してください。サポートされている地域のリストはアマゾンの岩盤地帯でご覧いただけます。サポートされているリージョンに切り替えるには、次のコマンドを実行します。
aws configure set region <your-region>TiDB Cloud Starterインスタンス
お持ちでない場合は、 TiDB Cloud Starterインスタンスを作成する。
Amazon Bedrockに必要な権限AWS アカウントと次のモデルへのアクセス:
- Amazon Titan Embeddings (
amazon.titan-embed-text-v2:0) は、テキスト埋め込みを生成するために使用されます。 - テキスト生成に使用されるメタラマ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接続情報を取得し、開発環境の環境変数を以下のように設定してください。
私のTiDBページに移動し、次に、対象のTiDB Cloud Starterインスタンスの名前をクリックして、概要ページに移動します。
右上隅の「接続」をクリックしてください。接続ダイアログが表示されます。
接続ダイアログの設定がご使用のオペレーティング環境と一致していることを確認してください。
接続タイプは
Publicに設定されています。ブランチは
mainに設定されています。Connect With は
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.