Azure SDK を使って Python から Azure AI Search のインデックスにjsonファイルのドキュメントを追加する

執筆日

2024/05/07

概要

「Azure AI Searchのインデックスを作成するpythonスクリプト」で作った検索インデックスにドキュメントを追加するスクリプトの書き方の紹介記事になります。
インデックス作成でドキュメント構成を定義したので、今回はその定義に従って作ったjsonファイルをドキュメントとして登録していきます。

前提

• python実行環境
• Azure AI Searchのリソースを作成済み

依存ライブラリインストール

pip install azure-search-documents

Note

• 各ドキュメントのid
  • ドキュメントを一意に示せるidであればintではなくstrでもいいですが特殊文字が使えないなのどの制約があるため、とりあえず識別番号を振り分けるのが安心です
  • PJで"{filename}_{page}"をidにしようとしたら……
    • 何かしらでエンコーディングすれば回避できる
• 既存インデックス
  • 既存インデックスに含まれるドキュメントの確認のため、スクリプト内に検索に関する処理を少し含みます
  • 既存ドキュメントを消したい場合のオプション引数を用意していますが「既存の数 > 新規の数」の場合に対応していないので注意してください
• 大量のドキュメントをアップロードする場合
  • 一度に大量のドキュメントをアップロードしようとするとRequestEntityTooLargeError: Operation returned an invalid status 'Request Entity Too Large'というエラーが出て怒られます
  • この記事のスクリプトには含まれませんが、PJでは3000以上のドキュメントを一気にアップロードしようとして怒られたので、500ずつにスライスしてアップロードするように書き換えました
• インデックスのフィールドについて
  • 追加ドキュメントのkeyがインデックスのフィールドと一致していないとエラーが出ます
  • 次元数が指定されているベクトルフィールドに空のリストを渡すのは大丈夫でした

スクリプト

既存インデックスへのドキュメント追加が出来るように序盤で既存インデックスの中身を調べる部分を少し長く書いてしまっていますが、やっていることはSearchClientを呼び出して、読み込んだjsonファイルに含まれる辞書形式のドキュメントを追加しているだけで単純です。

$ python upload_documents.py <--new> # --newオプションを付けると既存ドキュメントを消します

upload_documents.py

import os, json
from argparse import ArgumentParser
from dotenv import load_dotenv

from azure.core.credentials import AzureKeyCredential
from azure.search.documents import SearchClient

if __name__ == "__main__":
    parser = ArgumentParser(description="インデックスの更新か新規作成の選択ができます")
    parser.add_argument("-n", "--new", action="store_true", help="今あるドキュメントをリセットしてドキュメントを追加する")
    args = parser.parse_args()
    new_index = args.new

    # Azure AI Searchのエンドポイント・キー・および保存先インデックス名
    load_dotenv("../environment/.env")
    azure_search_endpoint = os.getenv("AZURE_SEARCH_SERVICE_ENDPOINT")
    azure_search_key = os.getenv("AZURE_SEARCH_ADMIN_KEY")
    azure_search_index_name = os.getenv("AZURE_SEARCH_INDEX")
    
    # 対象のインデックス名を指定してクライアントを作成  
    client = SearchClient(azure_search_endpoint, azure_search_index_name, AzureKeyCredential(azure_search_key))

    # 追加するドキュメントの開始idを、現在あるドキュメント数に設定
    # オプションでnewを設定した場合はid=0からにして上書きしていくが完全にリセットするように書いていないので注意
    if not new_index:
        # すべてのドキュメントが検索されるように*検索
        results = client.search(search_text="*", include_total_count=True) 
        # 既存ドキュメントと被らないようにするならタイトルリストを作るなどする
        # document_name_list = [result["title"] for result in results]
        # 既存のドキュメント数から追加するドキュメントのidを決めて割り振る
        id = results.get_count()
    else:
        id = 0

    json_docs_path = "documents.json" # json形式でドキュメントのリストになっているファイルを指定
    with open(json_docs_path, "r", encoding="utf-8") as file:
        documents = json.load(file)
        for document in documents:
            document["id"] = id
            id += 1
    result = client.upload_documents(documents=new_documents)

• documents.jsonの例
  • 「Azure AI Searchのインデックスを作成するpythonスクリプト」で作ったインデックスの形式を採用するのであればこのような例になります
  • 1ドキュメントの区切りをチャンクと言います
    • どこからどこまでを1チャンクにするのがいいかは資料の性質によります
    • 会話ログなら質疑応答の1セットにしたり、パワポならページ毎にしたり、長い文章ならトークン数で区切ったりなど様々な区切りが考えられます

documents.jsonの例

[
    {
        "title": "<ドキュメント_1のタイトル>",
        "category": "<ドキュメント_1のカテゴリ>",
        "content": "<ドキュメント_1の内容>",
        "keywords": ["key_1", "key_2", ...],
        "content_vector": [<contentを埋め込んで作った1536次元ベクトル>]
    },
    {
        "title": "<ドキュメント_2のタイトル>",
        "category": "<ドキュメント_2のカテゴリ>",
        "content": "<ドキュメント_2の内容>",
        "keywords": ["key_2", "key_3", ...],
        "content_vector": [<contentを埋め込んで作った1536次元ベクトル>]
    },
    {...},
]