LlamaIndex Azure Blob Storageからベクトルデータを取得する

こんばんは、本日はLlamaIndexをプロダクト開発に利用するときに、使える技を紹介したいと思います。

タイトルの通りベクトルデータをAzure Blob Storageから取得してロードしたいと思います。

1 背景
2 ストレージアカウントにてコンテナを作成
3 ソースコード
4 コード解説
5 まとめ

背景

プロダクト環境でLlamaIndexを利用する際は、インデックスを定期的に更新する方法とインデックスの格納場所などを考慮する必要があります。

今回はインデックスをAzureに格納するケースを想定して、クラウド上のデータをローカルに落とすことなく読み込む方法を紹介します。

Azure Functionなど内部にファイルを書き込めないシステムの場合、インデックスファイルをローカルにダウンロードできないので、この手法が必要となります。

ストレージアカウントにてコンテナを作成

ストレージアカウント作成後、コンテナを作成しインデックスファイルを格納しておきます。

ソースコード

以下コードは必要な部分の抜粋です。

import os

from llama_index.storage.storage_context import StorageContext
from llama_index.storage.docstore import SimpleDocumentStore
from llama_index.vector_stores import SimpleVectorStore
from llama_index.storage.index_store import SimpleIndexStore
from llama_index import load_index_from_storage

from azure.storage.blob import BlobServiceClient

def load_index_from_blob():
	connect_str = os.environ['AZURE_STORAGE_CONNECTION_STRING']
	blob_service_client = BlobServiceClient.from_connection_string(connect_str)
	container_client = blob_service_client.get_container_client("<YOUR_CONTAINER_NAME>")

  blob_content = container_client.download_blob("storage/docstore.json").readall()
  docstore_json = json.loads(blob_content)
  print("docstore_json loaded from blob storage")
    
  blob_content = container_client.download_blob("storage/vector_store.json").readall()
  vector_store_json = json.loads(blob_content)
  print("vector_store_json loaded from blob storage")
  
  blob_content = container_client.download_blob("storage/index_store.json").readall()
  index_store_json = json.loads(blob_content)
  print("index_store_json loaded from blob storage")

	storage_context = StorageContext.from_defaults(
      docstore=SimpleDocumentStore.from_dict(docstore_json),
      vector_store=SimpleVectorStore.from_dict(vector_store_json),
      index_store=SimpleIndexStore.from_dict(index_store_json),
  )

	app_index = load_index_from_storage(storage_context)
  return app_index

コード解説

接続文字列を用いてAzure Blob Storageの通信クライアントを初期化

connect_str = os.environ['AZURE_STORAGE_CONNECTION_STRING']
blob_service_client = BlobServiceClient.from_connection_string(connect_str)
container_client = blob_service_client.get_container_client("<YOUR_CONTAINER_NAME>")

Blob StorageからファイルをダウンロードしJSON変換

blob_content = container_client.download_blob("storage/docstore.json").readall()
docstore_json = json.loads(blob_content)

JSONデータをLlamaIndexのインデックスファイルとしてロード

from_dictを使用することでインデックスファイルからファイル読み込みするのではなく、JSON形式のデータを直接利用することができます。

（結構、from_persist_dirを使ったサンプルコードはネット上で見ますが、from_dictを使ってるケースはあまりありませんでしたので参考にしてください。）

storage_context = StorageContext.from_defaults(
    docstore=SimpleDocumentStore.from_dict(docstore_json),
    vector_store=SimpleVectorStore.from_dict(vector_store_json),
    index_store=SimpleIndexStore.from_dict(index_store_json),
)