問題意識とやってみたこと

LangChain v0.352で、LangChain-Communityパッケージに属するAzure AI Searchのvectorstoreモジュールを使用するには、azure-search-documentsのv11.4.0（最新のstable版）ではなくv11.4.0b8という古いβバージョンを使用しないといけない。そのため、Langchainのモジュールを通してAzure AI Searchの操作を行うときに、少し古いβバージョーンのAPI仕様を意識しないといけない（ベクトル検索の設定など）。
わざわざ古い仕様を覚えるのは嫌なので、LangChain v0.352のAzure AI Searchのvectorstoreモジュールをazure-search-documentsのv11.4.0で動くようにクラスを継承してカスタムクラスを作る。

継承して作ったカスタムクラスのコード

全量

CustomAzureSearchクラスの定義

import os
from google.colab import userdata

from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores.azuresearch import AzureSearch

from azure.core.credentials import AzureKeyCredential
from azure.core.exceptions import ResourceNotFoundError
from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential
from azure.search.documents import SearchClient
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.models import VectorizedQuery
from azure.search.documents.indexes.models import (
    SearchIndex,
    ScoringProfile,
    SearchableField,
    SearchField,
    SearchFieldDataType,
    SimpleField,
    TextWeights,
    CorsOptions,
)
from azure.search.documents.indexes.models import (
    VectorSearch,
    VectorSearchProfile,
    HnswAlgorithmConfiguration,  # HnswVectorSearchAlgorithmConfiguration is old
    SemanticPrioritizedFields,  # PrioritizedFields outdated
    SemanticSearch,  # SemanticSettings outdated
    SemanticConfiguration,
    SemanticField,
)

FIELDS_ID = "id"
FIELDS_CONTENT = "content"
FIELDS_CONTENT_VECTOR = "content_vector"
FIELDS_METADATA = "metadata"
MAX_UPLOAD_BATCH_SIZE = 1000

import numpy as np
import base64
import json
import uuid
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Iterable,
    List,
    Optional,
    Tuple,
    Type,
    Union,
)
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings

def _get_search_client(
    endpoint: str,
    key: str,
    index_name: str,
    semantic_configuration_name: Optional[str] = None,
    fields: Optional[List[SearchField]] = None,
    vector_search: Optional[VectorSearch] = None,
    semantic_search: Optional[SemanticSearch] = None,
    scoring_profiles: Optional[List[ScoringProfile]] = None,
    default_scoring_profile: Optional[str] = None,
    default_fields: Optional[List[SearchField]] = None,
    user_agent: Optional[str] = "langchain",
    cors_options: Optional[CorsOptions] = None,
) -> SearchClient:

    default_fields = default_fields or []
    if key is None:
        credential = DefaultAzureCredential()
    elif key.upper() == "INTERACTIVE":
        credential = InteractiveBrowserCredential()
        credential.get_token("https://search.azure.com/.default")
    else:
        credential = AzureKeyCredential(key)
    index_client: SearchIndexClient = SearchIndexClient(
        endpoint=endpoint, credential=credential, user_agent=user_agent
    )
    try:
        index_client.get_index(name=index_name)
    except ResourceNotFoundError:
        # Fields configuration
        if fields is not None:
            # Check mandatory fields
            fields_types = {f.name: f.type for f in fields}
            mandatory_fields = {df.name: df.type for df in default_fields}
            # Check for missing keys
            missing_fields = {
                key: mandatory_fields[key]
                for key, value in set(mandatory_fields.items())
                - set(fields_types.items())
            }
            if len(missing_fields) > 0:
                # Helper for formatting field information for each missing field.
                def fmt_err(x: str) -> str:
                    return (
                        f"{x} current type: '{fields_types.get(x, 'MISSING')}'. "
                        f"It has to be '{mandatory_fields.get(x)}' or you can point "
                        f"to a different '{mandatory_fields.get(x)}' field name by "
                        f"using the env variable 'AZURESEARCH_FIELDS_{x.upper()}'"
                    )

                error = "\n".join([fmt_err(x) for x in missing_fields])
                raise ValueError(
                    f"You need to specify at least the following fields "
                    f"{missing_fields} or provide alternative field names in the env "
                    f"variables.\n\n{error}"
                )
        else:
            fields = default_fields
        # Vector search configuration
        if vector_search is None:

            # >= 11.4.0:
            #   VectorSearch(algorithm_configuration) --> VectorSearch(profiles, algorithms)
            # HnswVectorSearchAlgorithmConfiguration --> HnswAlgorithmConfiguration
            vector_search = VectorSearch(
                profiles=[
                    VectorSearchProfile(
                        name="default",
                        algorithm_configuration_name="default",
                    )
                ],
                algorithms=[
                    HnswAlgorithmConfiguration(
                        name="default",
                        kind="hnsw",
                        parameters={  # type: ignore
                            "m": 4,
                            "efConstruction": 400,
                            "efSearch": 500,
                            "metric": "cosine",
                        },
                    )
                ]
            )

        # Create the semantic settings with the configuration
        if semantic_search is None and semantic_configuration_name is not None:

            # <=11.4.0: SemanticSettings --> SemanticSearch
            # PrioritizedFields(prioritized_content_fields)
            #   --> SemanticPrioritizedFields(content_fields)
            semantic_search = SemanticSearch(
                configurations=[
                    SemanticConfiguration(
                        name=semantic_configuration_name,
                        prioritized_fields=SemanticPrioritizedFields(
                            content_fields=[
                                SemanticField(field_name=FIELDS_CONTENT)
                            ],
                        ),
                    )
                ]
            )

        # Create the search index with the semantic settings and vector search
        index = SearchIndex(
            name=index_name,
            fields=fields,
            vector_search=vector_search,
            semantic_search=semantic_search,
            scoring_profiles=scoring_profiles,
            default_scoring_profile=default_scoring_profile,
            cors_options=cors_options,
        )
        index_client.create_index(index)
    # Create the search client
    return SearchClient(
        endpoint=endpoint,
        index_name=index_name,
        credential=credential,
        user_agent=user_agent,
    )

class CustomAzureSearch(AzureSearch):
    def __init__(
        self,
        azure_search_endpoint: str,
        azure_search_key: str,
        index_name: str,
        embedding_function: Callable,
        search_type: str = "hybrid",
        semantic_configuration_name: Optional[str] = None,
        semantic_query_language: str = "en-us",
        fields: Optional[List[SearchField]] = None,
        vector_search: Optional[VectorSearch] = None,
        semantic_search: Optional[SemanticSearch] = None,
        scoring_profiles: Optional[List[ScoringProfile]] = None,
        default_scoring_profile: Optional[str] = None,
        cors_options: Optional[CorsOptions] = None,
        **kwargs: Any,
    ):
        """Initialize with necessary components."""
        # Initialize base class
        self.embedding_function = embedding_function
        default_fields = [
            SimpleField(
                name=FIELDS_ID,
                type=SearchFieldDataType.String,
                key=True,
                filterable=True,
            ),
            SearchableField(
                name=FIELDS_CONTENT,
                type=SearchFieldDataType.String,
            ),
            SearchField(
                name=FIELDS_CONTENT_VECTOR,
                type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
                searchable=True,
                vector_search_dimensions=len(embedding_function("Text")),
                vector_search_profile_name="default",
            ),
            SearchableField(
                name=FIELDS_METADATA,
                type=SearchFieldDataType.String,
            ),
        ]
        user_agent = "langchain"
        if "user_agent" in kwargs and kwargs["user_agent"]:
            user_agent += " " + kwargs["user_agent"]
        self.client = _get_search_client(
            azure_search_endpoint,
            azure_search_key,
            index_name,
            semantic_configuration_name=semantic_configuration_name,
            fields=fields,
            vector_search=vector_search,
            semantic_search=semantic_search,
            scoring_profiles=scoring_profiles,
            default_scoring_profile=default_scoring_profile,
            default_fields=default_fields,
            user_agent=user_agent,
            cors_options=cors_options,
        )
        self.search_type = search_type
        self.semantic_configuration_name = semantic_configuration_name
        self.semantic_query_language = semantic_query_language
        self.fields = fields if fields else default_fields

    def vector_search_with_score(
        self, query: str, k: int = 4, filters: Optional[str] = None
    ) -> List[Tuple[Document, float]]:
        """Return docs most similar to query.

        Args:
            query: Text to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.

        Returns:
            List of Documents most similar to the query and score for each
        """
        results = self.client.search(
            search_text="",
            vector_queries=[
                VectorizedQuery(
                    vector=np.array(
                        self.embedding_function(query), dtype=np.float32
                    ).tolist(),
                    k_nearest_neighbors=k,
                    fields=FIELDS_CONTENT_VECTOR,
                )
            ],
            filter=filters,
        )
        # Convert results to Document objects
        docs = [
            (
                Document(
                    page_content=result.pop(FIELDS_CONTENT),
                    metadata={
                        **(
                            {FIELDS_ID: result.pop(FIELDS_ID)}
                            if FIELDS_ID in result
                            else {}
                        ),
                        **(
                            json.loads(result[FIELDS_METADATA])
                            if FIELDS_METADATA in result
                            else {
                                k: v
                                for k, v in result.items()
                                if k != FIELDS_CONTENT_VECTOR
                            }
                        ),
                    },
                ),
                float(result["@search.score"]),
            )
            for result in results
        ]
        return docs

    def hybrid_search_with_score(
        self, query: str, k: int = 4, filters: Optional[str] = None
    ) -> List[Tuple[Document, float]]:
        """Return docs most similar to query with an hybrid query.

        Args:
            query: Text to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.

        Returns:
            List of Documents most similar to the query and score for each
        """
        results = self.client.search(
            search_text=query,
            vector_queries=[
                VectorizedQuery(
                    vector=np.array(
                        self.embedding_function(query), dtype=np.float32
                    ).tolist(),
                    k_nearest_neighbors=k,
                    fields=FIELDS_CONTENT_VECTOR,
                )
            ],
            filter=filters,
            top=k,
        )
        # Convert results to Document objects
        docs = [
            (
                Document(
                    page_content=result.pop(FIELDS_CONTENT),
                    metadata={
                        **(
                            {FIELDS_ID: result.pop(FIELDS_ID)}
                            if FIELDS_ID in result
                            else {}
                        ),
                        **(
                            json.loads(result[FIELDS_METADATA])
                            if FIELDS_METADATA in result
                            else {
                                k: v
                                for k, v in result.items()
                                if k != FIELDS_CONTENT_VECTOR
                            }
                        ),
                    },
                ),
                float(result["@search.score"]),
            )
            for result in results
        ]
        return docs

    def semantic_hybrid_search_with_score_and_rerank(
        self, query: str, k: int = 4, filters: Optional[str] = None
    ) -> List[Tuple[Document, float, float]]:
        """Return docs most similar to query with an hybrid query.

        Args:
            query: Text to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.

        Returns:
            List of Documents most similar to the query and score for each
        """
        results = self.client.search(
            search_text=query,
            vector_queries=[
                VectorizedQuery(
                    vector=np.array(
                        self.embedding_function(query), dtype=np.float32
                    ).tolist(),
                    k_nearest_neighbors=50,
                    fields=FIELDS_CONTENT_VECTOR,
                )
            ],
            filter=filters,
            query_type="semantic",
            query_language=self.semantic_query_language,
            semantic_configuration_name=self.semantic_configuration_name,
            query_caption="extractive",
            query_answer="extractive",
            top=k,
        )
        # Get Semantic Answers
        semantic_answers = results.get_answers() or []
        semantic_answers_dict: Dict = {}
        for semantic_answer in semantic_answers:
            semantic_answers_dict[semantic_answer.key] = {
                "text": semantic_answer.text,
                "highlights": semantic_answer.highlights,
            }
        # Convert results to Document objects
        docs = [
            (
                Document(
                    page_content=result.pop(FIELDS_CONTENT),
                    metadata={
                        **(
                            {FIELDS_ID: result.pop(FIELDS_ID)}
                            if FIELDS_ID in result
                            else {}
                        ),
                        **(
                            json.loads(result[FIELDS_METADATA])
                            if FIELDS_METADATA in result
                            else {
                                k: v
                                for k, v in result.items()
                                if k != FIELDS_CONTENT_VECTOR
                            }
                        ),
                        **{
                            "captions": {
                                "text": result.get("@search.captions", [{}])[0].text,
                                "highlights": result.get("@search.captions", [{}])[
                                    0
                                ].highlights,
                            }
                            if result.get("@search.captions")
                            else {},
                            "answers": semantic_answers_dict.get(
                                json.loads(result["metadata"]).get("key"), ""
                            ),
                        },
                    },
                ),
                float(result["@search.score"]),
                float(result["@search.reranker_score"]),
            )
            for result in results
        ]
        return docs

主要な変更点

主な変化点は、ベクトル検索を行うFieldの設定方法と、VectorSearchの定義方法と、ベクトル検索の仕方にあります。すべてベクトル検索周りですね。

VectorSearchの定義方法（algorithmsだけでなくprofilesも引数に取るようになった）

vector_search = VectorSearch(
    profiles=[
        VectorSearchProfile(
            name="default",
            algorithm_configuration_name="default",
        )
    ],
    algorithms=[
        HnswAlgorithmConfiguration(
            name="default",
            kind="hnsw",
            parameters={  # type: ignore
                "m": 4,
                "efConstruction": 400,
                "efSearch": 500,
                "metric": "cosine",
            },
        )
    ]
)

ベクトル検索を行うFieldの設定方法（vector_search_profile_nameを新たな引数に持つようになった。これには、VectorSearchの定義で引数に渡したprofileの名前を渡す）

SearchField(
    name=FIELDS_CONTENT_VECTOR,
    type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
    searchable=True,
    vector_search_dimensions=len(embedding_function("Text")),
    vector_search_profile_name="default",
),

ベクトル検索の仕方（searchメソッドの引数の書き方が変わった）

results = self.client.search(
    search_text=query,
    vector_queries=[
        VectorizedQuery(
            vector=np.array(
                self.embedding_function(query), dtype=np.float32
            ).tolist(),
            k_nearest_neighbors=k,
            fields=FIELDS_CONTENT_VECTOR,
        )
    ],
    filter=filters,
    top=k,
)

動作確認

LangChainのAzure AI Search VectorSotreのチュートリアルを参考に動作確認を行います。

OpenAIやAzure AI SearchのAPIKEYやエンドポイントの設定

Google colabolartoryのシークレットという機能を利用しています（userdata.getのところ）。

os.environ["OPENAI_API_KEY"] = userdata.get('OPENAI_API_BASE')
model: str = "text-embedding-ada-002"
vector_store_address: str = userdata.get('AZURE_SEARCH_SERVICE_ENDPOINT')
vector_store_password: str = userdata.get('AZURE_SEARCH_API_KEY')

インデックスの作成

embeddings: OpenAIEmbeddings = OpenAIEmbeddings(deployment=model, chunk_size=1)
embedding_function = embeddings.embed_query


fields = [
    SimpleField(
        name="id",
        type=SearchFieldDataType.String,
        key=True,
        filterable=True,
    ),
    SearchableField(
        name="content",
        type=SearchFieldDataType.String,
        searchable=True,
    ),
    SearchField(
        name="content_vector",
        type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
        searchable=True,
        vector_search_dimensions=len(embedding_function("Text")),
        vector_search_profile_name="default",
    ),
    SearchableField(
        name="metadata",
        type=SearchFieldDataType.String,
        searchable=True,
    ),
    # Additional field to store the title
    SearchableField(
        name="title",
        type=SearchFieldDataType.String,
        searchable=True,
    ),
    # Additional field for filtering on document source
    SimpleField(
        name="source",
        type=SearchFieldDataType.String,
        filterable=True,
    ),
]

index_name: str = "langchain-vector-demo-custom"

vector_store: AzureSearch = CustomAzureSearch(
    azure_search_endpoint=vector_store_address,
    azure_search_key=vector_store_password,
    index_name=index_name,
    embedding_function=embedding_function,
    fields=fields,
)

ドキュメントのアップロード

# Data in the metadata dictionary with a corresponding field in the index will be added to the index
# In this example, the metadata dictionary contains a title, a source and a random field
# The title and the source will be added to the index as separate fields, but the random won't. (as it is not defined in the fields list)
# The random field will be only stored in the metadata field
vector_store.add_texts(
    ["Test 1", "Test 2", "Test 3"],
    [
        {"title": "Title 1", "source": "A", "random": "10290"},
        {"title": "Title 2", "source": "A", "random": "48392"},
        {"title": "Title 3", "source": "B", "random": "32893"},
    ],
)

検索

例1 ハイブリッド検索

res = vector_store.similarity_search(query="Test 3 source1", k=3, search_type="hybrid")
res

ハイブリッド検索の出力

[Document(page_content='Test 3', metadata={'id': 'NWE2OWJjNTQtMDM3MS00ZDg5LTgwMGYtN2UwMzNhZjhmNjlh', 'title': 'Title 3', 'source': 'B', 'random': '32893'}),
 Document(page_content='Test 1', metadata={'id': 'MTlkYmM3MTMtZDYyZS00Y2ExLThmMzctYTY2NGM2ODFiMmVh', 'title': 'Title 1', 'source': 'A', 'random': '10290'}),
 Document(page_content='Test 2', metadata={'id': 'N2RjNDc1YmEtZjc2MC00N2YwLTljNzItYmRlZDcyMWJkODk4', 'title': 'Title 2', 'source': 'A', 'random': '48392'})]

例2 フィルターありハイブリッド検索

フィルターありハイブリッド検索

res = vector_store.similarity_search(
    query="Test 3 source1", k=3, search_type="hybrid", filters="source eq 'A'"
)
res

フィルターありハイブリッド検索の出力

[Document(page_content='Test 1', metadata={'id': 'MTlkYmM3MTMtZDYyZS00Y2ExLThmMzctYTY2NGM2ODFiMmVh', 'title': 'Title 1', 'source': 'A', 'random': '10290'}),
 Document(page_content='Test 2', metadata={'id': 'N2RjNDc1YmEtZjc2MC00N2YwLTljNzItYmRlZDcyMWJkODk4', 'title': 'Title 2', 'source': 'A', 'random': '48392'})]

まとめと所感

最新のazure-search-moduleとLangChainとカスタムしたコードで、LangChainの公式のチュートリアルのコードが動くところまで確かめました。Semantic検索の部分は動かしていないのでもしかしたらうまくいかないかもです。
以下感想です。
LangChainのAzure AI SearchのVectorStoreの実装を見てみて、Azure AI Searchの機能、特に全文検索の機能をを活かしきれていないなと感じました。検索エンジンではなくて、VectorStoreなのだから、全文検索の機能を活かせなてなくてもその名に偽りがないとは思います。しかし、それでは他のベクトルDBに比べてAzure AI Searchを使う旨みが、機能面ではSemantic検索やハイブリッド検索を簡単に実施できることぐらいしか無さそう、という印象を持ちました。
Azure AI Searchの機能をもっと活かすなら、既存のLangChainの実装を参考に独自のVectorStoreを作ってもよいかなと思いました。個人的には、Documentクラスの扱い方など参考にしたいです。