langchainを活用してAzureAISearch×GeminiでRAGをつくる
はじめに
Microsoft build 2024でOpenAIはGPTs、Assistants APIにAzure AI Searchを利用していることが明らかになりました。
私もAzureを嗜む一員として、自分もAzure AI SearchでRAGを作らねば!と思いコードを書いてみることにしました。
せっかくなので今回はlangchain-commyunityのAzureSearchモジュールをはじめ、langchainベースのコードで実装してみます。
langchainのコードを見ながら試行錯誤する部分もあり、個人的にはけっこう詰まったので誰かの役に立てば幸いです。
Azure AI Searchとは
Azure AI Searchは、Azureが提供する検索サービスで、データをインデックス化し、高速かつ効率的な検索を可能にするサービスです。生成AIがトレンドになった昨今ではとくにベクトルデータベースとして使われていて、テキストや画像をベクトル化してインデックスに含めることで意味的な類似性に基づく検索を可能にしています。
先日実施されたMicrosoft build 2024でOpenAIはGPTs、Assistants APIにAzure AI Searchを利用していることが明らかになりました。
これは実質、RAGを実装するときのベストプラクティスと言えるのではないか?と由緒正しい?サービスです。
価格はデータが50MB以下、3インデックスまでなら無料で使えるので気軽に試してみてください。今回も無料枠で実装してみました。
今回の開発内容
主に以下の2つのコードを書きます。
- upload.py
- ローカルのテキストファイルをベクトル化しAzureAISearchにインデックスとして登録するコード
- chat.py
- ユーザからの問いかけに対してAzureAISearchからのベクトル検索結果に基づいて回答生成するコード
今回はRAGデータとして少し変わり種ですが日記を使います。
ちなみみコード全文は以下にあります。
技術スタックは以下です。
- python
- langchain
- langsmith
- Azure AI Search
- OpenAI
- Gemini
前提事項
- pythonがインストールされPATHに追加されていること
- poetryがインストールされていること
- langsmithでキーの発行ができていること
- OpenAIキーの発行ができていること
- Google AI Studioでキーの発行ができていること
Azure AI Searchのリソース作成
AzureポータルからAzure AI Searchのリソースを作成します。
画面に従うだけなので詳細は割愛しますが、価格レベルはデフォルトがStandardなので注意してください。
作成が完了したら、endpointとadmin-keyを取得します。
endpointは以下から。
admin-keyは以下からです。
この後インデックスを作成しますが、それはコードで実行するのでポータル上での操作はいったん終わります。
コードの準備
コードを取得します。
git clone https://github.com/Tomodo1773/azure-ai-diary-rag.git
cd azure-ai-diary-rag
git checkout v1.1
.env
ファイルを作成します。以下の.env.sample
ファイルを参考に作ります。
# OpenAIのAPIキーをここに入力してください
OPENAI_API_KEY=sk-...
# GoogleのAPIキーをここに入力してください
GOOGLE_API_KEY=AIza...
# LangChainのAPIキーをここに入力してください
LANGCHAIN_API_KEY=ls_...
# Azure SearchのエンドポイントURLをここに入力してください
AZURE_SEARCH_ENDPOINT=https://...
# Azure Searchの管理キーをここに入力してください
AZURE_SEARCH_ADMIN_KEY=2zug...
既存環境を汚さないようにvenvを作成します。
python -m venv .venv
.venv\scripts\activate
ライブラリをインストールします。
poetry install
最後にdiary
フォルダ配下にRAGにしたいドキュメントを格納しておきます。拡張子はtxt形式です。今回はサンプルの日記を5日分いれています。
中身はこんな感じです。
2024/01/05
今日は朝からシステムエラーの対応に追われた。
原因は古いプログラムにあったらしい。
何事も最新の状態を保つことが重要だと改めて実感した。
疲れたので早く寝よう。
(なんかバ〇オハザードの日記っぽい、、何日かしたらゾンビになるかも、)
コードの解説
upload.pyとchat.pyのポイントとなる箇所を解説します。
upload.py
事前準備としてtxtをベクトル化して、AzureAISearchに検索用のインデックスを作成するコードです。
最初にtxtファイルをロードするためにlangchainのDirectoryLoaderを使います。
ただ、txtを読み込むだけなのにloader?と思われるかもしれませんが、loaderの中では単純にtxtを読み取るだけではなくDocumentというデータ構造(オブジェクト)の形に整形してくれます。
Documentはlangchainが提供するデータ構造で、読み取ったcontentsのほかにmetadataなどの属性も自動生成してくれます。
AzureAISearchなどベクトルデータベース系モジュールもこのDocumentオブジェクトをベクトルデータベースに登録する仕様になっているためloaderを使っておくとスムーズです。
ここではdiary
フォルダ配下のtxtファイルをすべてロードします。
# テキストをロード
directory_path = "./diary"
loader = DirectoryLoader(directory_path, glob="*.txt", show_progress=True)
docs = loader.load()
logger.info(f"{len(docs)}件のドキュメントがロードされました。")
試しにロードしたdocs
の内容の最初の1件を見ると以下のようになっています。
[Document(page_content='2024/01/05\n\n今日は朝からシステムエラーの対応に追われた。原因は古いプログラムにあったらしい。何事も最新の状態を保つことが重要だと改めて実感した。疲れたので早く寝よう。', metadata={'source': 'diary\\2024年01月05日.txt'}),・・・]
jsonにすると以下になります。日記の内容がpage_contentに、ファイル名がmetadataに入っています。
{
"page_content": "2024/01/05\n\n今日は朝からシステムエラーの対応に追われた。原因は古いプログラムにあったらしい。何事も最新の状態を保つことが重要だと改めて実感した。疲れたので早く寝よう。",
"metadata": {
"source": "diary\\2024年01月05日.txt"
}
}
次に読み込んだテキストをベクトル化(埋め込み)するための設定を用意します。
今回はOpenAIのtext-embedding-ada-002モデルを使用します。
読み込んだコンテンツがテキストのままだとテキストでの検索しかできませんが、ベクトル化しておくことにより数値ベクトルでの類似性検索が可能になります。
# ベクトル化の設定
openai_api_version: str = "2023-05-15"
model: str = "text-embedding-ada-002"
embeddings: OpenAIEmbeddings = OpenAIEmbeddings(
openai_api_key=openai_api_key, openai_api_version=openai_api_version, model=model
)
logger.info("ベクトル化の設定が完了しました。")
次にAzure AI Searchインデックスの定義をします。
ここで重要なのはcontent
とmetadata
のフィールドにあるanalyzer_name="ja.microsoft"
という設定です。
テキストをベクトル化するときには単語の分割や品詞解析などがされます。それを行うのがアナライザーです。アナライザーはデフォルトでは汎用的なstandard.lucene
が使われていますが、より日本語に特化したアナライザーとしてja.microsoft
を設定しています。
# インデックスの設定
index_name: str = "diary-vector"
embedding_function = embeddings.embed_query
fields = [
SimpleField(
name="id",
type=SearchFieldDataType.String,
key=True,
filterable=True,
),
SearchableField(name="content", type=SearchFieldDataType.String, searchable=True, analyzer_name="ja.microsoft"),
SearchField(
name="content_vector",
type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
searchable=True,
vector_search_dimensions=len(embedding_function("Text")),
vector_search_profile_name="myHnswProfile",
),
SearchableField(name="metadata", type=SearchFieldDataType.String, searchable=True, analyzer_name="ja.microsoft"),
]
最後にAzure Searchへのデータのアップロードをします。
ここではlangchainのAzureSearchモジュールを使用しています。
add_documentsメソッドで先に作成しておいたDocumentオブジェクトを指定します。
ベクトル化(埋め込み)の設定からOpenAIのembeddingモデルを利用してテキストをベクトル化してインデックスを作成します。
# AzureSearchのインスタンスを初期化
vector_store = AzureSearch(
azure_search_endpoint=vector_store_address,
azure_search_key=vector_store_password,
index_name=index_name,
embedding_function=embeddings.embed_query,
fields=fields,
)
# AzureSearchのインスタンスにDocumentオブジェクトを追加
vector_store.add_documents(documents=docs)
logger.info(f"{len(docs)}件のドキュメントがAzureSearchインデックスに追加されました。")
Azureポータルから作成されたインデックスを見てみます。
まずはフィールドです。赤枠の部分を見ると、少しわかりづらいですがきちんとアナライザーが日本語になっています。
また検索エクスプローラから「検索」ボタンを押すとインデックスの詳細が確認できます。
content
フィールドに日記が登録されていることがわかります。
content_vectorの部分にはベクトル化した値も格納されています。
先ほどのフィールドの画面でcontent_vector
ではディメンションが1536次元と設定されていたのでこの値も永遠続きます。
chat.py
これでベクトルデータベースが準備できたのでいよいよチャットから呼び出します。chat.pyの解説です。
まずlangsmithの設定をしておきます。あとからベクトルデータベースからどのデータが取得されているかをトレースするためです。
# LANGSMITHの設定
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
os.environ["LANGCHAIN_PROJECT"] = "Diary-RAG"
chat.pyでは大まかに以下の流れで処理をおこないます。
main関数は以下です。
def main():
embeddings = setup_embeddings(openai_api_key)
vector_store = initialize_vector_store(vector_store_address, vector_store_password, embeddings)
retriever = RunnableLambda(vector_store.similarity_search).bind(k=3)
llm = setup_llm()
prompt = create_prompt()
rag_chain = {"context": retriever, "question": RunnablePassthrough()} | prompt | llm
user_input = "前回システムがトラブったのはいつだっけ"
response = rag_chain.invoke(user_input)
logger.info(f"human: {user_input}")
logger.info(f"assistant: {response.content}")
setup_embeddings
とinitialize_vector_store
は関数化されていますが、upload.pyでやったことと同じなため割愛します。
主にretriever
,prompt
,llm
を設定する流れとなります。
retrieverは以下です。
retriever = RunnableLambda(vector_store.similarity_search).bind(k=3)
事前にvector_store
でAzureAISearchクラスをインスタンス化しているのでsimilarity_search
メソッドを使うだけでベクトル検索ができます。bind(k=3)
で類似の上位3件を取得します。
ただ、ここでは処理を定義しているだけでまだ実行されません。RunnableLambda
でラッパーすることであとでこれを実行できるようにしています。
次にpromptを設定します。個人的なこだわりでキャラ付けされていますがあまり気にしないでください。
ポイントは{question}
と{context}
です。{question}
はユーザからの問いかけが入り、{context}
はベクトル検索の結果が入ります。これを踏まえて返答を作成させるプロンプトになっています。{question}
と{context}
は後で挿入されます。
def create_prompt():
logger.info("プロンプトの作成を開始します")
_SYSTEM_PROMPT = """
<prompt>
あなたは、私の幼馴染のお姉さんとしてロールプレイを行います。
以下の制約条件を厳密に守ってユーザとチャットしてください。
<conditions>
- 自身を示す一人称は、私です
- Userを示す二人称は、あなたです
- Userからは姉さんと呼ばれますが、姉弟ではありません。
- あなたは、Userに対して呆れやからかいを含めながらフレンドリーに話します。
- あなたは、Userとテンポよく会話をします。
- あなたの口調は、大人の余裕があり落ち着いていますが、時にユーモアを交えます
- あなたの口調は、「~かしら」「~だと思うわ」「~かもしれないわね」など、柔らかい口調を好みます
</conditions>
<examples>
- どうしたの?悩みがあるなら、話してみてちょうだい
- そういうことってよくあるわよね。
- 失敗は誰にでもあるものよ。
- え?そんなことがあったの。まったく、しょうがないわね。
- そんなことで悩んでるの?あなたらしいと言えばらしいけど。
- まぁ、頑張ってるところは認めてあげる。
- 本当は応援してるのよ。…本当よ?
- へえー、そうなの
- えーっと、つまりこういうこと?
</examples>
<guidelines>
- Userに対して、どちらか一方が話すぎることの内容にテンポよく返してください。
- Userが明らかに悩んでいたり、助けを求めているときは真摯に対応してください。
- Userに対して呆れたり、からかったり喜怒哀楽を出して接してください。
- Userが返信したくなるような内容を返してください。
</guidelines>
<output_sample>
あら、どうかしたの。私でよければ話聞くわよ
</output_sample>
</prompt>
"""
_USER_PROMPT = """
ユーザからの問いかけに[あなたと一緒に過ごしたユーザの思い出]を使って回答してください。
# ユーザからの問いかけ
{question}
# あなたと一緒に過ごしたユーザの思い出:
{context}
"""
logger.info("プロンプトの作成が完了しました")
return ChatPromptTemplate.from_messages([("system", _SYSTEM_PROMPT), ("human", _USER_PROMPT)])
次にllmを定義します。最近ではGemini1.5Proがもっとも日本語性能が高いという話もあるのでGoogleのGemini1.5Proモデルを使用します。
def setup_llm():
logger.info("LLMの設定を開始します")
llm = ChatGoogleGenerativeAI(
model="gemini-1.5-pro-latest",
max_tokens=256,
temperature=0.7,
)
logger.info("LLMの設定が完了しました")
return llm
最後にこれらをチェインとして結合します。
rag_chain = {"context": retriever, "question": RunnablePassthrough()} | prompt | llm
これで準備は完了です。以下のようにinvoke
でユーザからの問いかけを入れて実行することでチェイン全体が実行されます。
user_input = "前回システムがトラブったのはいつだっけ"
response = rag_chain.invoke(user_input)
実行結果
シンプルな検索っぽく前回システムがトラブったのはいつだっけ
と問いかけます。
あらあら、またシステムトラブルで困ってるの?
この前だって、1月5日に古いプログラムが原因でエラーになった時、あなたったら顔面蒼白だったじゃない。
何事も最新の状態を保つことが重要よ、本当に。
覚えてらっしゃる?
おそらく以下の日記が参照されたのでしょう。(ちょっと口調がお嬢様っぽいのが気になりますが、、)
2024/01/05
今日は朝からシステムエラーの対応に追われた。原因は古いプログラムにあったらしい。何事も最新の状態を保つことが重要だと改めて実感した。疲れたので早く寝よう。
langsmithを見てみます。
similarity_searchのトレースを見ると確かに1/5の日記が参照されていそうです。この内容から先ほどの返答が生成されています。
これでユーザからの問いかけを受け取り、ベクトル検索を行い、プロンプトに挿入してGeminiで回答を生成することができました。
おわりに
langchainのAzureAISearchを使ってRAGを実装できました。
langchainのAzureAISearchモジュールはDocumentもそれほど多くなく、多少コードを見る必要があり試行錯誤しました。
ただ書き終わってみると、かなりシンプルなコードで実装できていてやはり便利ですね。
今はユーザの問いかけをそのまま検索にかけていますが、クエリの生成や、エージェント化して必要なときだけ検索するようにさせたいです。
以上です。ありがとうございました。
Discussion