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

このチュートリアルでは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以降を推奨）。

ヒント

完全なサンプルコードJupyter Notebook で表示することも、 コラボレーションオンライン環境で直接実行することもできます。

前提条件

このチュートリアルを完了するには、以下が必要です。

• Python 3.8以降インストールされています。
• ジュピターノートブックがインストールされました。
• Gitがインストールされました。
• TiDBクラスタ。

TiDBクラスタをお持ちでない場合は、以下の手順で作成できます。

• (推奨) TiDB Cloud Starterインスタンスを作成する。
• ローカルテスト用のTiDBセルフマネージドクラスタをデプロイ。または本番本番のTiDBセルフマネージドクラスタをデプロイ。

さあ始めましょう

このセクションでは、TiDB Vector SearchをLlamaIndexと統合してセマンティック検索を実行するための手順を段階的に説明します。

ステップ1. 新しいJupyter Notebookファイルを作成する

ルートディレクトリに、 integrate_with_llamaindex.ipynbという名前の新しい Jupyter Notebook ファイルを作成します。

touch integrate_with_llamaindex.ipynb

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

プロジェクトディレクトリで、以下のコマンドを実行して必要なパッケージをインストールしてください。

pip install llama-index-vector-stores-tidbvector
pip install llama-index

Jupyter Notebookでintegrate_with_llamaindex.ipynbファイルを開き、必要なパッケージをインポートするために以下のコードを追加します。

import textwrap

from llama_index.core import SimpleDirectoryReader, StorageContext
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.tidbvector import TiDBVectorStore

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

選択したTiDBのデプロイオプションに応じて、環境変数を設定してください。

# Use getpass to securely prompt for environment variables in your terminal.
import getpass
import os

# Connection string format: "mysql+pymysql://<USER>:<PASSWORD>@<HOST>:4000/<DB>?ssl_ca=/etc/ssl/cert.pem&ssl_verify_cert=true&ssl_verify_identity=true"
tidb_connection_string = getpass.getpass("TiDB Connection String:")
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")

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

ステップ4. サンプルドキュメントを読み込む

ステップ4.1 サンプル文書をダウンロードする

プロジェクトディレクトリ内に、 data/paul_graham/という名前のディレクトリを作成し、 run-llama/llama_index GitHubリポジトリからサンプルドキュメントpaul_graham_essay.txtをダウンロードしてください。

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

ステップ4.2 ドキュメントを読み込む

data/paul_graham/paul_graham_essay.txtクラスを使用して、 SimpleDirectoryReader } からサンプル ドキュメントを読み込みます。

documents = SimpleDirectoryReader("./data/paul_graham").load_data()
print("Document ID:", documents[0].doc_id)

for index, document in enumerate(documents):
   document.metadata = {"book": "paul_graham"}

ステップ5．ドキュメントベクターを埋め込んで保存する

ステップ5.1 TiDBベクターストアの初期化

以下のコードは、ベクトル検索に最適化されたTiDBにpaul_graham_testという名前のテーブルを作成します。

tidbvec = TiDBVectorStore(
   connection_string=tidb_connection_string,
   table_name="paul_graham_test",
   distance_strategy="cosine",
   vector_dimension=1536,
   drop_existing_table=False,
)

正常に実行されると、TiDB データベース内のpaul_graham_testテーブルを直接表示およびアクセスできるようになります。

ステップ5.2 埋め込みを生成して保存する

以下のコードは、ドキュメントを解析し、埋め込みを生成し、それらをTiDBベクトルストアに保存します。

storage_context = StorageContext.from_defaults(vector_store=tidbvec)
index = VectorStoreIndex.from_documents(
   documents, storage_context=storage_context, show_progress=True
)

期待される出力は以下のとおりです。

Parsing nodes: 100%|██████████| 1/1 [00:00<00:00,  8.76it/s]
Generating embeddings: 100%|██████████| 21/21 [00:02<00:00,  8.22it/s]

ステップ6. ベクトル検索を実行する

以下では、TiDBベクトルストアに基づいてクエリエンジンを作成し、意味的類似性検索を実行します。

query_engine = index.as_query_engine()
response = query_engine.query("What did the author do?")
print(textwrap.fill(str(response), 100))

注記

TiDBVectorStore defaultクエリ モードのみをサポートしています。

期待される出力は以下のとおりです。

The author worked on writing, programming, building microcomputers, giving talks at conferences,
publishing essays online, developing spam filters, painting, hosting dinner parties, and purchasing
a building for office use.

ステップ7．メタデータフィルターを使用して検索する

検索結果を絞り込むには、メタデータフィルターを使用して、適用したフィルターに一致する特定の近隣検索結果を取得できます。

book != "paul_graham"フィルターと異なるクエリ

次の例ではbookメタデータ フィールドが"paul_graham"である結果を除外しています。

from llama_index.core.vector_stores.types import (
   MetadataFilter,
   MetadataFilters,
)

query_engine = index.as_query_engine(
   filters=MetadataFilters(
      filters=[
         MetadataFilter(key="book", value="paul_graham", operator="!="),
      ]
   ),
   similarity_top_k=2,
)
response = query_engine.query("What did the author learn?")
print(textwrap.fill(str(response), 100))

期待される出力は以下のとおりです。

Empty Response

book == "paul_graham"フィルターを含むクエリ

次の例ではbookメタデータ フィールドが"paul_graham"であるドキュメントのみを結果に含めるようにフィルタリングしています。

from llama_index.core.vector_stores.types import (
   MetadataFilter,
   MetadataFilters,
)

query_engine = index.as_query_engine(
   filters=MetadataFilters(
      filters=[
         MetadataFilter(key="book", value="paul_graham", operator="=="),
      ]
   ),
   similarity_top_k=2,
)
response = query_engine.query("What did the author learn?")
print(textwrap.fill(str(response), 100))

期待される出力は以下のとおりです。

The author learned programming on an IBM 1401 using an early version of Fortran in 9th grade, then
later transitioned to working with microcomputers like the TRS-80 and Apple II. Additionally, the
author studied philosophy in college but found it unfulfilling, leading to a switch to studying AI.
Later on, the author attended art school in both the US and Italy, where they observed a lack of
substantial teaching in the painting department.

ステップ8．文書を削除する

インデックスから最初のドキュメントを削除します。

tidbvec.delete(documents[0].doc_id)

文書が削除されているかどうかを確認してください。

query_engine = index.as_query_engine()
response = query_engine.query("What did the author learn?")
print(textwrap.fill(str(response), 100))

期待される出力は以下のとおりです。

Empty Response

関連項目

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