SmartHR Tech Blog

SmartHR 開発者ブログ

PythonでAzure AI Searchのインデックスを構築する方法

PythonでAzure AI Searchのインデックスを構築する方法
SmartHRでAIプロダクトの開発をしているmizunaoです!現在SmartHRでは、LLM(Large Language Model)を活用したプロダクト開発に取り組んでいます。この記事では、その一環として取り組んだ「PythonでAzure AI Searchのインデックスを構築する方法」を紹介します。

Azure AI Searchとは?

Azure AI SearchはAzureが提供する検索サービスで、全文検索、ベクトル検索、ハイブリッド検索など、様々な検索ソリューションに対応しています。検索に使用するファイルはPDFやOffice系ファイルなど、様々な形式に対応しています。 特に、RAG(Retrieval-Augmented Generation)において、検索の前段階のコンテキスト取得にAzure AI Searchは広く使われており、社内文書やヘルプページなど、LLMが持たない情報を検索したい時に非常に有効な検索基盤となっています。

SmartHRでは、インフラは主にGoogle Cloudを利用しています。そのため、Vertex AI SearchやVertex AI Agent Builderも検討しましたが、最終的には使いやすさやカスタマイズ性を重視し、Azure AI Searchを採用することに決めました。

必要な前準備

インデックスを構築するには、事前に次の2つのステップを完了しておく必要があります。

1. Azureリソースのセットアップ

Azure Portalで以下のリソースを作成してください。

2. 必要なPythonパッケージのインストール

PythonからAzure AI Searchを操作するにはazure-search-documentsというAzure SDKをインストールする必要があります。今回はバージョン11.6.0b4を使用します。

pip install azure-search-documents==11.6.0b4

インデックスの構築

インデックスを構築するにあたって、主なステップは次の通りです。

  1. インデックスの作成
  2. データソースの作成
  3. スキルセットの定義
  4. インデクサーの作成・実行

まず、Azure AI Searchのリソースを作成・更新するために必要なインデクサーのクライアントを生成しておきます。

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

client = SearchIndexerClient(
    "your-azure-search-service-endpoint", AzureKeyCredential("your-azure-search-api-key")
)

1. インデックスの作成

インデックスは、検索を高速化するために構造化されたデータストレージです。全文検索、ベクトル検索、ハイブリッド検索などの検索に対応しています。検索対象に含めたいフィールドはsearchable=True を指定する必要があります。chunkにチャンク化された文字列、text_vectorにベクトル化されたデータが保存されます。

from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex,
    SearchField,
    SearchFieldDataType,
    VectorSearch,
    VectorSearchProfile,
    HnswAlgorithmConfiguration,
    SemanticSearch,
    SemanticConfiguration,
    SemanticPrioritizedFields,
    SemanticField
)

def _create_index() -> SearchIndex:
    index = SearchIndex(
        name="your-index-name",
        fields=[
            SearchField(
                name="chunk_id",
                type=SearchFieldDataType.String,
                key=True,
                hidden=False,
                filterable=True,
                sortable=True,
                facetable=True,
                searchable=True,
                analyzer_name="keyword",
            ),
            SearchField(
                name="chunk",
                type=SearchFieldDataType.String,
                hidden=False,
                filterable=False,
                sortable=False,
                facetable=False,
                searchable=True,
                analyzer_name="ja.lucene",
            ),
            SearchField(
                name="title",
                type=SearchFieldDataType.String,
                hidden=False,
                filterable=True,
                sortable=False,
                facetable=False,
                searchable=True,
                analyzer_name="ja.lucene",
            ),
            SearchField(
                name="text_vector",
                type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
                hidden=True,
                filterable=False,
                sortable=False,
                facetable=False,
                searchable=True,
                vector_search_dimensions=3072,
                vector_search_profile_name="profile",
            ),
        ],
        # ベクトルプロファイル
        vector_search=VectorSearch(
            profiles=[
                VectorSearchProfile(
                    name="profile",
                    algorithm_configuration_name="hnsw-algorithm",
                )
            ],
            algorithms=[HnswAlgorithmConfiguration(name="hnsw-algorithm")],
        ),
        # セマンティック構成
        semantic_search=SemanticSearch(
            configurations=[
                SemanticConfiguration(
                    name="your-semantic-config-name",
                    prioritized_fields=SemanticPrioritizedFields(
                        title_field=SemanticField(field_name="title"),
                        content_fields=[SemanticField(field_name="chunk")],
                    ),
                )
            ]
        ),
    )
    index_client = SearchIndexClient(
        "your-azure-search-service-endpoint", AzureKeyCredential("your-azure-search-api-key")
    )
    index_client.create_or_update_index(index)

    return None

2. データソースの設定

Azure Blob Storageをデータソースの接続先として指定します。これにより、ストレージ内のデータをインデクサーが読み取れるようになります。

from azure.search.documents.indexes.models import SearchIndexerDataSourceConnection, SearchIndexerDataContainer

def _create_datasource() -> None:
    data_source_connection = SearchIndexerDataSourceConnection(
        name="your-datasource-name",
        type="azureblob",
        connection_string="your-connection-string",
        container=SearchIndexerDataContainer(name="your-container-name"),
    )
    client.create_or_update_data_source_connection(data_source_connection)

    return None

3. スキルセットの作成

Azure AI Searchはデータをベクトル化する際にスキルを適用させることができます。スキルにはOCRやチャンク分割など、様々なスキルがあります。ここでは以下の4つのスキルを使用してファイルデータをベクトルに変換しています。

  1. OCR Skillでファイルからテキストを抽出
  2. Merge Skillで抽出したテキストや他のテキストソースを結合する
  3. Split Skillで結合されたテキストをチャンクに分割
  4. Embedding Skillでチャンクをベクトルに変換
from azure.search.documents.indexes.models import (
    InputFieldMappingEntry,
    OutputFieldMappingEntry,
    OcrSkill,
    SplitSkill,
    MergeSkill,
    AzureOpenAIEmbeddingSkill,
    SearchIndexerSkillset,
)

def _create_skillset() -> None:
    ocr_skill = OcrSkill(
        description="OCR Skill",
        context="/document/normalized_images/*",
        default_language_code="ja",
        inputs=[
            InputFieldMappingEntry(
                name="image", source="/document/normalized_images/*"
            )
        ],
        outputs=[OutputFieldMappingEntry(name="text", target_name="text")],
    )

    merge_skill = MergeSkill(
        description="Merge Skill",
        context="/document",
        inputs=[
            InputFieldMappingEntry(name="text", source="/document/content"),
            InputFieldMappingEntry(
                name="itemsToInsert", source="/document/normalized_images/*/text"
            ),
            InputFieldMappingEntry(
                name="offsets", source="/document/normalized_images/*/contentOffset"
            ),
        ],
        outputs=[
            OutputFieldMappingEntry(name="mergedText", target_name="extracted_text")
        ],
    )

    split_skill = SplitSkill(
        description="Split Skill",
        inputs=[InputFieldMappingEntry(name="text", source="/document/extracted_text")],
        outputs=[OutputFieldMappingEntry(name="textItems", target_name="pages")],
        maximum_page_length=2048,
        page_overlap_length=512,
        maximum_pages_to_take=0,
        default_language_code="ja",
    )

    embedding_skill = AzureOpenAIEmbeddingSkill(
        description="Embedding Skill",
        inputs=[InputFieldMappingEntry(name="text", source="/document/pages/*")],
        outputs=[OutputFieldMappingEntry(name="embedding", target_name="text_vector")],
        context="/document/pages/*",
        resource_uri="your-azure-openai-endpoint",
        api_key="your-azure-openai-key",
        deployment_id="text-embedding-3-large",
        model_name="text-embedding-3-large",
        dimensions=3072,
    )

    skillset = SearchIndexerSkillset(
        name="your-skillset-name",
        skills=[ocr_skill, merge_skill, split_skill, embedding_skill]
    )
    client.create_or_update_skillset(skillset)

    return None

4. インデクサーの作成・実行

インデクサーは、Azure Blob Storageに接続しているデータソースのファイルデータをスキルセットに基づいて変換処理を行い、インデックスに格納してくれるクローラーです。

from azure.search.documents.indexes.models import (
    IndexingParametersConfiguration,
    IndexingParameters,
    SearchIndexer,
    FieldMapping,
)

def _create_indexer() -> None:
    configuration = IndexingParametersConfiguration(
        data_to_extract="contentAndMetadata",
        parsing_mode="default",
        image_action="generateNormalizedImages",
        query_timeout=None,
        pdf_text_rotation_algorithm="detectAngles",
    )

    indexer = SearchIndexer(
        name="your-indexer-name",
        target_index_name="your-index-name",
        data_source_name="your-data-source-name",
        skillset_name="your-skillset-name",
        parameters=IndexingParameters(
            configuration=configuration,
        ),
        field_mappings=[
            FieldMapping(
                source_field_name="metadata_storage_name",
                target_field_name="title",
            ),
        ],
    )

    client.create_or_update_indexer(indexer)

    return None

実行方法

上記のコードをもとに、全体を統合して実行します。実行するとAzure Blob Storageで指定したコンテナ内のファイルがベクトルに変換され、インデックスに保存されます。

def create_ai_search() -> None:
    _create_index()
    _create_datasource()
    _create_skillset()
    _create_indexer()

    print("インデックス作成・インデクサー実行完了")
    return None

create_ai_search()

結果の確認

Azure Portalからインデックスとインデクサーのステータスを確認することができます。また、search_client.search("検索クエリ")を実行するとインデックス内のデータを全文検索することができます。ベクトル検索する場合は、検索クエリをtext-embedding-3-largeなどのモデルでベクトル化してから検索してください。

from azure.search.documents import SearchClient

search_client = SearchClient(
    "your-azure-search-service-endpoint",
    "your-index-name",
    AzureKeyCredential("your-azure-search-api-key"),
)

results = search_client.search("検索クエリ")
for result in results:
    print(result)

インデックスを構築する時にしたほうがよいこと

インデックスの作成とインデクサーの実行が完了すれば、検索が可能になります。しかし、運用をスムーズに進めるためには、事前に考慮しておくべきポイントがいくつかあるので紹介します。

インデックスのフィールドにはチャンク元のファイルを特定するIDを追加しておくべき

Azure AI Searchで検索すると、データはチャンクごとに分割された状態で取得されます。そのため、検索結果から元のファイルを特定するためのファイルIDをインデックスに追加しておくことをおすすめします。これにより、検索結果から元の文書を辿る際や、特定のファイルに基づいたデータ分析を行う際に役立ちます。

# file_idを追加する場合、Azure Blob Storageに保存しているファイルのメタデータにfile_idを追加しておく必要があります。
SearchField(
    name="file_id",
    type=SearchFieldDataType.String,
    hidden=False,
    filterable=True,
    sortable=False,
    facetable=False,
    searchable=True,
)

変換処理でエラーが出た際の最大失敗数を設定しておくべき

インデクサーがファイルデータをベクトルデータに変換する際、エラーが発生することがあります。デフォルトの設定では、1つでもエラーが発生すると全体の処理が停止してしまいます。そのため、最大失敗数(max_failed_items)を設定して、ある程度エラーが発生するまで変換処理を止めないようにしておいた方が良いかもしれません。

parameters=IndexingParameters(
    max_failed_items=3
)

まとめ

PythonからAzure AI Searchのインデックスを構築してみました。RAGを構築する際にAzure AI Searchは非常に有用なのでぜひ試してみてください。

Yes, we’re hiring!

いかがでしたでしょうか。SmartHRのAI専任チームでは以下の職種を募集しています!まずはカジュアル面談からでも大歓迎です。ご応募お待ちしております!