ベクトル検索をLlamaIndexと統合する
このチュートリアルでは、TiDB Cloudのベクトル検索機能をラマインデックスと統合する方法を説明します。
注記
- TiDB Vector Search は現在ベータ版であり、 TiDB Cloudサーバーレスクラスターでのみ使用できます。
- 完全なサンプルコード Jupyter Notebook で表示したり、サンプル コードをコラボオンライン環境で直接実行したりできます。
前提条件
このチュートリアルを完了するには、次のものが必要です。
- Python 3.8以上インストールされました。
- ジュピターノートブックインストールされました。
- ギットインストールされました。
- TiDB Cloud Serverless クラスター。TiDB Cloud クラスターがない場合は、 TiDB Cloud Serverless クラスターの作成に従って独自のTiDB Cloudクラスターを作成してください。
始める
このセクションでは、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. 環境を設定する
ステップ3.1 TiDBクラスタへの接続文字列を取得する
クラスターページに移動し、ターゲット クラスターの名前をクリックして概要ページに移動します。
右上隅の「接続」をクリックします。接続ダイアログが表示されます。
接続ダイアログの構成が動作環境と一致していることを確認します。
- 接続タイプは
Publicに設定されています。 - ブランチは
mainに設定されています。 - Connect With は
SQLAlchemyに設定されています。 - オペレーティング システムは環境に適合します。
- 接続タイプは
PyMySQLタブをクリックし、接続文字列をコピーします。
ヒント：
まだパスワードを設定していない場合は、 「パスワードの生成」をクリックしてランダムなパスワードを生成します。
ステップ3.2 環境変数を設定する
安全で効率的なデータベース接続を確立するには、 TiDB Cloudが提供する標準の接続方法を使用します。
このドキュメントでは、埋め込みモデル プロバイダーとしてオープンAI使用します。この手順では、手順 3.1 から取得した接続文字列とOpenAI APIキーを指定する必要があります。
環境変数を設定するには、次のコードを実行します。接続文字列と OpenAI API キーを入力するよう求められます。
import getpass
import os
tidb_connection_url = getpass.getpass(
"TiDB connection URL (format - mysql+pymysql://root@127.0.0.1:4000/test): "
)
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
ステップ4. サンプル文書を読み込む
ステップ4.1 サンプルドキュメントをダウンロードする
プロジェクト ディレクトリにdata/paul_graham/という名前のディレクトリを作成し、 ランラマ/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 ドキュメントを読み込む
SimpleDirectoryReaderクラスを使用してdata/paul_graham/paul_graham_essay.txtからサンプル ドキュメントを読み込みます。
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ベクターストアを初期化する
次のコードは、ベクトル検索に最適化されたpaul_graham_testという名前のテーブルを TiDB に作成します。
tidbvec = TiDBVectorStore(
connection_string=tidb_connection_url,
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))
注記
TiDBVectorStoredefaultクエリ モードのみをサポートします。
期待される出力は次のとおりです。
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