Azure AI Search のシノニムマップをPythonで作成してみた

"""
Azure AI Searchのシノニムマップを作成・更新するスクリプト
"""

import json
import logging
import os
import sys
from datetime import datetime
from typing import Dict, List, Optional

from azure.core.credentials import AzureKeyCredential
from azure.core.exceptions import AzureError, ResourceNotFoundError
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import SearchField, SearchIndex, SynonymMap
from pydantic import BaseModel, Field

# ログ設定
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler("synonym_update.log"),
        logging.StreamHandler(sys.stdout),
    ],
)
logger = logging.getLogger(__name__)

class SynonymConfig(BaseModel):
    """Azure AI Searchシノニム設定"""

    # Azure AI Search接続設定
    service_endpoint: str = Field(..., description="AI Search service endpoint")
    admin_key: str = Field(..., description="AI Search admin key for write operations")
    index_name: str = Field(..., description="AI Search index name")

    # シノニム設定
    synonym_file_path: str = Field(..., description="シノニム定義JSONファイルのパス")
    synonym_map_name: str = Field(
        default="business-terms", description="シノニムマップ名（JSONファイル内のnameフィールドを優先）"
    )

    # 実行設定
    dry_run: bool = Field(default=False, description="実際には更新せずにログ出力のみ")
    overwrite_existing: bool = Field(
        default=True, description="既存のシノニムマップを上書きするか"
    )

class SynonymUpdater:
    """Azure AI Searchのシノニムマップを更新するクラス"""

    def __init__(self, config: SynonymConfig):
        self.config = config
        self.index_client = SearchIndexClient(
            endpoint=config.service_endpoint,
            credential=AzureKeyCredential(config.admin_key),
        )

    def _load_synonym_definition(self) -> Dict:
        """シノニム定義JSONファイルを読み込む"""
        try:
            with open(self.config.synonym_file_path, "r", encoding="utf-8") as f:
                data = json.load(f)

            logger.info(f"Loaded synonym definition from: {self.config.synonym_file_path}")
            logger.debug(f"Synonym data: {data}")

            # 必須フィールドの確認
            if "synonyms" not in data:
                raise ValueError("Missing 'synonyms' field in JSON file")

            # nameフィールドがあれば使用、なければconfigのデフォルトを使用
            if "name" in data:
                self.config.synonym_map_name = data["name"]
                logger.info(f"Using synonym map name from JSON: {self.config.synonym_map_name}")

            return data

        except FileNotFoundError:
            logger.error(f"Synonym file not found: {self.config.synonym_file_path}")
            raise
        except json.JSONDecodeError as e:
            logger.error(f"Invalid JSON in synonym file: {e}")
            raise
        except Exception as e:
            logger.error(f"Error loading synonym definition: {e}")
            raise

    def _create_synonym_map(self, synonym_data: Dict) -> SynonymMap:
        """シノニムマップオブジェクトを作成"""
        try:
            # シノニムデータを適切な形式に変換
            synonyms_raw = synonym_data["synonyms"]

            if isinstance(synonyms_raw, str):
                # 文字列の場合は改行で分割してリストに変換
                synonyms_list = [line.strip() for line in synonyms_raw.split('\n') if line.strip()]
                logger.info(f"Converted synonym string to list: {len(synonyms_list)} rules")
                logger.debug(f"Synonym rules: {synonyms_list}")
            elif isinstance(synonyms_raw, list):
                # すでにリストの場合はそのまま使用
                synonyms_list = synonyms_raw
                logger.info(f"Using synonym list as-is: {len(synonyms_list)} rules")
            else:
                raise ValueError(f"Invalid synonym format: {type(synonyms_raw)}")

            # SynonymMapオブジェクトを作成
            synonym_map = SynonymMap(
                name=self.config.synonym_map_name,
                synonyms=synonyms_list,
                encryption_key=synonym_data.get("encryptionKey"),  # オプション
            )

            logger.info(f"Created synonym map object: {synonym_map.name}")
            return synonym_map

        except Exception as e:
            logger.error(f"Error creating synonym map object: {e}")
            raise

    def _check_existing_synonym_map(self) -> Optional[SynonymMap]:
        """既存のシノニムマップを確認"""
        try:
            existing_map = self.index_client.get_synonym_map(self.config.synonym_map_name)
            logger.info(f"Found existing synonym map: {existing_map.name}")
            return existing_map
        except ResourceNotFoundError:
            logger.info(f"No existing synonym map found: {self.config.synonym_map_name}")
            return None
        except Exception as e:
            logger.error(f"Error checking existing synonym map: {e}")
            raise

    def _update_index_fields(self) -> None:
        """インデックスのフィールドにシノニムマップを適用"""
        try:
            # 現在のインデックス定義を取得
            index = self.index_client.get_index(self.config.index_name)
            logger.info(f"Retrieved index: {index.name}")

            # インデックス内の全フィールドをデバッグ出力
            logger.info("Available fields in index:")
            for field in index.fields:
                field_type = field.type.name if hasattr(field.type, 'name') else str(field.type)
                logger.info(f"  Field: {field.name}, Type: {field_type}, Searchable: {field.searchable}")

            # シノニムマップを適用するフィールドを特定
            # chunk と title フィールドのみに適用
            updated_fields = []
            fields_to_update = []

            for field in index.fields:
                field_type = field.type.name if hasattr(field.type, 'name') else str(field.type)
                if (
                    field.searchable
                    and field_type == "Edm.String"
                    and field.name in ["chunk", "title"]  # 対象フィールドを限定
                ):
                    # synonym_map_namesプロパティを確認（Azure AI Searchの実際のプロパティ名）
                    field_dict = field.as_dict()
                    current_maps = field_dict.get("synonym_map_names", []) or []
                    logger.info(f"Field '{field.name}' current synonym maps: {current_maps}")

                    if self.config.synonym_map_name not in current_maps:
                        # フィールドのコピーを作成して更新
                        field_dict["synonym_map_names"] = [self.config.synonym_map_name]
                        updated_field = SearchField.from_dict(field_dict)
                        updated_fields.append(updated_field)
                        fields_to_update.append(field.name)
                        logger.info(f"Will update field '{field.name}' with synonym map '{self.config.synonym_map_name}'")
                    else:
                        updated_fields.append(field)
                        logger.info(f"Field '{field.name}' already has synonym map '{self.config.synonym_map_name}'")
                else:
                    updated_fields.append(field)
                    field_type = field.type.name if hasattr(field.type, 'name') else str(field.type)
                    if field.searchable and field_type == "Edm.String":
                        logger.info(f"Skipping field '{field.name}' (not in target list)")

            logger.info(f"Found {len(fields_to_update)} fields to update: {fields_to_update}")

            if not fields_to_update:
                logger.info("No fields need synonym map update")
                return

            if self.config.dry_run:
                logger.info(f"DRY RUN: Would update fields: {fields_to_update}")
                return

            # インデックスを更新
            logger.info("Updating index with synonym maps...")
            index.fields = updated_fields
            result = self.index_client.create_or_update_index(index)
            logger.info(f"Successfully updated index with synonym map on fields: {fields_to_update}")
            logger.info(f"Index update result: ETag={getattr(result, 'e_tag', 'N/A')}")

        except Exception as e:
            logger.error(f"Error updating index fields: {e}")
            logger.error(f"Error type: {type(e).__name__}")
            import traceback
            logger.error(f"Traceback: {traceback.format_exc()}")
            raise

    def run(self) -> None:
        """シノニムマップ更新処理を実行"""
        logger.info("Starting synonym map update process")
        logger.info(f"Service endpoint: {self.config.service_endpoint}")
        logger.info(f"Index name: {self.config.index_name}")
        logger.info(f"Synonym file: {self.config.synonym_file_path}")
        logger.info(f"Dry run: {self.config.dry_run}")

        try:
            # シノニム定義を読み込む
            synonym_data = self._load_synonym_definition()

            # シノニムマップオブジェクトを作成
            synonym_map = self._create_synonym_map(synonym_data)

            # 既存のシノニムマップを確認
            existing_map = self._check_existing_synonym_map()

            # シノニムマップを作成または更新
            if existing_map:
                if not self.config.overwrite_existing:
                    logger.warning(
                        f"Synonym map '{self.config.synonym_map_name}' already exists and overwrite_existing=False"
                    )
                    return

                if self.config.dry_run:
                    logger.info(
                        f"DRY RUN: Would update existing synonym map '{self.config.synonym_map_name}'"
                    )
                else:
                    self.index_client.create_or_update_synonym_map(synonym_map)
                    logger.info(f"Updated synonym map: {self.config.synonym_map_name}")
            else:
                if self.config.dry_run:
                    logger.info(f"DRY RUN: Would create new synonym map '{self.config.synonym_map_name}'")
                else:
                    self.index_client.create_synonym_map(synonym_map)
                    logger.info(f"Created new synonym map: {self.config.synonym_map_name}")

            # インデックスのフィールドにシノニムマップを適用
            logger.info("Applying synonym map to index fields...")
            self._update_index_fields()

            logger.info("Synonym map update process completed successfully")

        except Exception as e:
            logger.error(f"Synonym map update failed: {e}")
            raise

def main():
    """メイン実行関数"""

    # 環境変数から設定を読み込み
    service_endpoint = os.getenv("AZURE_SEARCH_SERVICE_ENDPOINT")
    admin_key = os.getenv("AZURE_SEARCH_ADMIN_KEY")
    index_name = os.getenv("AZURE_SEARCH_INDEX_NAME")
    synonym_file_path = os.getenv("SYNONYM_FILE_PATH", "synonyms.json")

    if not all([service_endpoint, admin_key, index_name]):
        logger.error("Required environment variables are missing:")
        logger.error("  AZURE_SEARCH_SERVICE_ENDPOINT")
        logger.error("  AZURE_SEARCH_ADMIN_KEY")
        logger.error("  AZURE_SEARCH_INDEX_NAME")
        sys.exit(1)

    # 設定を作成
    config = SynonymConfig(
        service_endpoint=service_endpoint,  # type: ignore
        admin_key=admin_key,  # type: ignore
        index_name=index_name,  # type: ignore
        synonym_file_path=synonym_file_path,
        dry_run=os.getenv("DRY_RUN", "false").lower() == "true",
        overwrite_existing=os.getenv("OVERWRITE_EXISTING", "true").lower() == "true",
    )

    # シノニムマップ更新実行
    updater = SynonymUpdater(config)
    updater.run()

if __name__ == "__main__":
    main()