執筆日

2024/05/02

概要

RAGなどの検索システムをAzure AI Searchで作成した際に、インデックス設定をテスト環境と運用環境で同一にするためにスクリプトで管理する方法があります。今回はそのインデックス作成のためのサンプルスクリプトを作りました。

前提

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

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

pip install azure-search-documents python-dotenv

Note

APIバージョンが変わると書き方が変わる場合が多々あるので注意してください。
Azure AI Searchについての記事は後日、

「インデックスにjsonファイルからドキュメントを追加する」
「インデックスに様々なクエリを投げる」
「Azure AI Searchのあれこれ（費用・各機能・Python SDKのドキュメントなど）」

等も書いていく予定です。

スクリプト

ベクトル検索やセマンティック検索など高度な検索を実装しない場合は、「インデックス構成要素作成」と「インデックス作成」のみ使用すれば大丈夫です。

インデックス構成要素作成

インデックスの構成要素を定義します

fieldsに加えたい要素を追加していきます
ベクトル検索・セマンティック検索・スコアリングプロファイルなどの高度な検索は必要に応じて実装します
- 通常の全文検索とベクトル検索を同時に行うことはハイブリッド検索と呼ばれ検索精度が向上することが知られています
- セマンティック検索は高度に意味的な分解を行うことで精度向上する手法でハイブリッド+セマンティック検索は非常に高い精度が期待できるようです

index_component.py (def create_index)

from azure.search.documents.indexes.models import (
    SearchIndex,
    SearchFieldDataType,
    # フィールド定義クラス
    SimpleField, # フィルタリングやソートに使用するフィールド
    SearchableField, # リッチな機能の全文検索を実装する場合
    SearchField, # 完全一致の全文検索を実装する場合
    # ベクトル検索実装のためのクラス
    VectorSearch,
    HnswAlgorithmConfiguration,
    VectorSearchProfile,
    # セマンティック検索実装のためのクラス
    SemanticConfiguration,
    SemanticPrioritizedFields,
    SemanticField,
    SemanticSearch,
    # スコアリングプロファイル実装のためのクラス
    ScoringProfile,
    TextWeights,
)
def create_rich_index_component(index_name: str, scoring_profiles_path = None, scoring_profile_name = None):
    # ベクトル検索とセマンティック検索を実装したインデックスの作成
    # DataSourceクラスを使えばDBからデータを取得することも可能
    fields = [
            SimpleField(name="id", type=SearchFieldDataType.String, key=True, sortable=True),
            SearchableField(name="title", type=SearchFieldDataType.String, analyzer="ja.microsoft", searchable=True),
            SearchField(name="category", type=SearchFieldDataType.String, filterable=True, facetable=True),
            SearchableField(name="content", type=SearchFieldDataType.String, analyzer="ja.microsoft", searchable=True),
            SearchField(name="keywords", type=SearchFieldDataType.Collection(SearchFieldDataType.String)), # SearchableFieldの方がいいかも
            # ベクトル検索を実装する場合
            SearchField(name="content_vector", type=SearchFieldDataType.Collection(SearchFieldDataType.Single), 
                        vector_search_dimensions=1536, vector_search_profile_name="vector_profile"), 
        ]

    # 高度な検索の実装
    vector_search = create_vector_search()
    semantic_search = create_semantic_search()
    scoring_profiles, default_scoring_profile = create_scoring_profiles(scoring_profiles_path, scoring_profile_name=scoring_profile_name)

    return SearchIndex(
        name=index_name,
        fields = fields,
        vector_search = vector_search, # 任意
        semantic_search = semantic_search, # 任意
        scoring_profiles = scoring_profiles, # 任意
        default_scoring_profile = default_scoring_profile, # 任意
    )

ベクトル検索, セマンティック検索の作成

ベクトル化したデータを検索対象にする場合やセマンティック検索を実装したい場合に使用します。

index_component.py (vector and semantic)

def create_vector_search(
        algorithm_name = "<任意のアルゴリズム名>", 
        vector_search_profile_name = "<任意のベクトルサーチプロファイル名>"
        ):
    # ベクトル検索を実装する場合
    vector_search = VectorSearch(
        profiles=[
            VectorSearchProfile(
                name=vector_search_profile_name,
                algorithm_configuration_name=algorithm_name,
            )
        ],
        algorithms=[HnswAlgorithmConfiguration(name=algorithm_name)],
    )

    return vector_search

def create_semantic_search(
        semantic_config_name="<任意のコンフィグ名>", 
        title_field_name="<タイトルに当たるフィールド名>", 
        content_field_name="<内容に当たるフィールド名>", 
        keywords_field_name="<キーワードに当たるフィールド名>"
        ):
    # セマンティック検索を実装する場合
    semantic_config = SemanticConfiguration(
        name=semantic_config_name,
        prioritized_fields=SemanticPrioritizedFields(
            title_field=SemanticField(field_name=title_field_name),
            content_fields=[SemanticField(field_name=content_field_name)],
            keywords_fields=[SemanticField(field_name=keywords_field_name)]
        )
    )
    semantic_search = SemanticSearch(configurations=[semantic_config])

    return semantic_search

スコアリングプロファイル作成

スコアリングプロファイルは同時に複数定義したり、重みのかけ方に複雑な関数を使ったりもできますが、一番簡単に1つのプロファイルを重み付け加算だけで実装した例です
- scoring_profiles変数に複数の重みをappendすることで複数のプロファイルを設定できます
- デフォルトのプロファイルを設定しないこともできますが、作ったプロファイルをデフォルトにする実装です
configファイルとして管理できるようにjsonファイルから重みを読み込む形にしています

scoring_profiles.json (タイトルに大きめの重みを付けてみた例)

{"title": 3, "content": 1, "keywords": 2}

index_component.py (def scoring_profile)

def create_scoring_profiles(
        scoring_profiles_path = None,
        function_aggregation = "sum", # "average"、"minimum"、"maximum"、"firstMatching"などもある
        scoring_profile_name = None,
        ):
    # スコアリングプロファイルの作成
    scoring_profiles = []
    if scoring_profiles_path:
        try:
            with open(scoring_profiles_path, "r") as f:
                weights = json.load(f)
            text_weights = TextWeights(weights=weights)
            scoring_profile = ScoringProfile(
                name=scoring_profile_name,
                function_aggregation=function_aggregation,
                text_weights=text_weights
            )
            scoring_profiles.append(scoring_profile)
            default_scoring_profile = scoring_profile_name
        except FileNotFoundError as e:
            print(e)
            print(f"{scoring_profiles_path} が見つかりませんでした。スコアリングプロファイルを作成できません。")
            scoring_profiles = None
            default_scoring_profile = None
    else:
        scoring_profiles = None
        default_scoring_profile = None

    return scoring_profiles, default_scoring_profile

インデックス作成

create_search_index.py

import os
from dotenv import load_dotenv

from azure.search.documents.indexes import SearchIndexClient
from azure.core.credentials import AzureKeyCredential

from index_component import create_index_component

load_dotenv("../environment/.env")
# Azure AI Searchのエンドポイント・キー・および保存先インデックス名
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 = SearchIndexClient(azure_search_endpoint, AzureKeyCredential(azure_search_key))
# インデックスの作成
index = create_index(index_name=azure_search_index_name,
    scoring_profiles_path = "<プロファイルの重みを記述したファイルのパス>",
    scoring_profile_name = "<任意のプロファイル名>")
result = client.create_or_update_index(index) # 指定したインデックス名が既存の場合上書きする
# result = client.create_index(index) # 上書きはしない

Discussion