【LangChain】Azure AI Searchを使ったカスタムRetrieverでRAGを自作する

3行まとめ

• LangChainでAI Searchを使ったRAGやりたい
• AI Search用のRetrieverはあるけど、セマンティック検索とかもできるようにしたい
• BaseRetrieverを使ってカスタマイズしてみる

AI Search用のRetrieverは既にあるけど

• セマンティック検索とかができない
• 弄れないパラメータも多い

やりたいこと

• retrieverを自作し、AI Searchはazureライブラリで使う
• retrieverに、queryと一緒にパラメータも投げられるようにする

  • パラメータのイメージ：

    inputs = {
        "query": "<query>", # ここにqueryを入れる
        "params": {
            "index_name": "<index_name>", # 検索するインデックス
            "search_key": "<search_key>", # 検索対象のキー
            "metadata_keys": [key1, key2], # metadataとして取得したいキーのリスト
            "top_k": n, # 何件検索結果を出すか
        }
    }
    

ソースコード

AI Searchを使ったRetriever

ai_search_retriever.py

from langchain.schema import BaseRetriever, Document
from azure.core.credentials import AzureKeyCredential
from azure.search.documents import SearchClient
from azure.search.documents.models import VectorizedQuery

class AISearchRetriever(BaseRetriever):
    """
    AzureAISearchRetriever
    """
    search_endpoint = "<your-ai-search-endpoint>"
    api_key = "<your-ai-search-api-key>"
    api_version = "<your-desired-api-version>"
    def _get_relevant_documents(self, inputs):
        documents = []
        query = inputs["query"]
        vector_query = model.encode(query)
        params = inputs["params"]
        search_key = params["search_key"]
        metadata_keys = params["metadata_keys"]
        select_keys = metadata_keys + [search_key]
        search_credential = AzureKeyCredential(self.api_key)
        search_client = SearchClient(
        endpoint=self.search_endpoint,
        index_name=params["index_name"],
        credential=search_credential,
        )
        results =  search_client.search(
            query_type="semantic",
            search_text=query,
            select=select_keys,
            vector_queries=[VectorizedQuery(
                vector=vector_query,
                k_nearest_neighbors=3, 
                fields=params["vector_key"],
            )],
            semantic_configuration_name="my-semantic-config",
            query_caption="extractive",
            top=params["top_k"],
            include_total_count=True,
        )
        for data in results:
            page_content = data.get(search_key)
            metadata = {key: value for key, value in data.items() if key in metadata_keys}
            doc = Document(page_content=page_content, metadata=metadata)
            documents.append(doc)
        return documents

AI Search Retrieverを使ったシンプルRAG

rag_with_ai_search.py

from langchain_openai import AzureChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

from ai_search_retriever import AISearchRetriever

import os
openai_endpoint = os.getenv("OPENAI_ENDPOINT")
openai_api_key = os.getenv("OPENAI_API_KEY")

ai_search_endpoint = os.getenv("AI_SEARCH_ENDPOINT")
ai_search_api_key = os.getenv("AI_SEARCH_API_KEY")

prompt_str = """
なんか入れる
{context}
"""

llm = AzureChatOpenAI(
    deployment_name="<deployment-name>",
    api_key=openai_api_key,
    azure_endpoint=openai_endpoint,
    openai_api_version="<your-desired-api-version>"
)

retriever = AISearchRetriever(
    search_endpoint="<your-ai-search-endpoint>",
    api_key="<your-ai-search-api-key>",
    api_version="<your-desired-api-version>"
)

prompt = ChatPromptTemplate.from_template(prompt_str)

rag_chain = {"context": retriever} | prompt | llm

inputs = {
    "query": "<query>",
    "params": {
        "index_name": "<index_name>",
        "search_key": "<search_key>",
        "metadata_keys": "<metadata_keys>",
        "top_k": "top_k",
    }
}

response = rag_chain.invoke(inputs)

print(response.content)

得られた知見

• retrieverに入れるinputって辞書でも(リストでも)よかったことを知らなかった
• BaseRetrieverを継承すればめちゃくちゃ自由にretrieverが作れる
  • 今回はAI Searchをazureライブラリを使って動かしているが、API叩くのもアリ
  • AI Searchに限らず何でもretrieverに使えそう