Spring AIでAzure Cosmos DBのベクトル検索を試す
概要
Azure Cosmos DBがベクトル検索できるようになり、Spring AIとも統合されました。
簡単なRAGを通じてSpring AIとAzure Cosmos DB、Azure OpenAI Serviceを用いたベクトル検索を試してみます。
コード例
今回のコード例は次のリポジトリのspring-ai-cosmos-vector-store-example
というディレクトリにあります。advent-calendar-20241225
というタグを参照してください。
以降の説明はこのコード例をもとにして記載しています。
バージョンなど
- macOS Sequoia 15.1.1 arm64
- Java 21
- Spring Boot 3.4.0
- Spring AI 1.0.0-M4
- azure-cli 2.67.0
Azureリソースの構築
構築するものは次の通りです。括弧内はリソースに付けた名前です。
- リソースグループ(
spring-ai-cosmos-vector-store-example
)- Azure Cosmos DB(
example-vector-store-<ランダム文字列>
) - Azure OpenAI Service(
example-openai-<ランダム文字列>
)- 埋め込みモデル(
example-embedding
)- モデルは
text-embedding-3-small
を利用
- モデルは
- チャットモデル(
example-chat
)- モデルは
gpt-4o-mini
を利用
- モデルは
- 埋め込みモデル(
- Azure Cosmos DB(
データベースやコンテナーはSpringが自動で作ってくれるのでここでは作りません。
各リソースはaz
コマンドとBicepファイルで構築します。
まずリソースグループを作成します。
resource_group="spring-ai-cosmos-vector-store-example"
az group create --name $resource_group
リソースグループが作成できたらBicepで残りのリソースを構築します。
# Azure Cosmos DBとAzure OpenAIのアカウント名を設定。重複を避けるため一部をランダム生成
cosmos_account_name="example-vector-store-$(cat /dev/urandom | LC_ALL=C tr -dc 'a-z0-9' | fold -w 10 | head -n 1)"
openai_account_name="example-openai-$(cat /dev/urandom | LC_ALL=C tr -dc 'a-z0-9' | fold -w 10 | head -n 1)"
# 埋め込みモデルとチャットモデルのデプロイメント名を設定
embedding_model_deployment_name="example-embedding"
chat_model_deployment_name="example-chat"
az deployment group create --resource-group $resource_group \
--template-file src/infra/vector_store.bicep \
--parameters cosmosAccountName=$cosmos_account_name \
openaiAccountName=$openai_account_name \
embeddingModelDeploymentName=$embedding_model_deployment_name \
chatModelDeploymentName=$chat_model_deployment_name
Spring Bootアプリケーションに読み込ませるため、いくつかの環境変数を設定しておきます。
# 埋め込みモデルとチャットモデルのデプロイメント名を環境変数へ設定する
export SPRING_AI_AZURE_OPENAI_EMBEDDING_OPTIONS_DEPLOYMENTNAME=$embedding_model_deployment_name
export SPRING_AI_AZURE_OPENAI_CHAT_OPTIONS_DEPLOYMENTNAME=$chat_model_deployment_name
# Azure Cosmos DBのエンドポイントとキーを取得して環境変数へ設定する
export SPRING_AI_VECTORSTORE_COSMOSDB_ENDPOINT=$(az cosmosdb show --resource-group $resource_group --name $cosmos_account_name --query "documentEndpoint" --output tsv)
export SPRING_AI_VECTORSTORE_COSMOSDB_KEY=$(az cosmosdb keys list --resource-group $resource_group --name $cosmos_account_name --query "primaryMasterKey" --output tsv)
# Azure OpenAIのエンドポイントとキーを取得して環境変数へ設定する
export SPRING_AI_AZURE_OPENAI_ENDPOINT=$(az cognitiveservices account show --resource-group $resource_group --name $openai_account_name --query "properties.endpoint" --output tsv)
export SPRING_AI_AZURE_OPENAI_APIKEY=$(az cognitiveservices account keys list --resource-group $resource_group --name $openai_account_name --query "key1" --output tsv)
Spring AIアプリケーションを実装する
基本的には冒頭にリンクした公式ドキュメントを参考に実装すればOKです。
ここでは以前に試したRAGのコードを流用します。
コントローラーのコードを掲載しておきます。
@RestController
public class RagController {
@Autowired
private ChatClient chatClient;
@Autowired
private VectorStore vectorStore;
@PostMapping("/load-docs")
public void loadDocs() {
vectorStore.add(List.of(new Document("harrison worked at kensho")));
}
@PostMapping("/rag")
public Object rag(@RequestParam String question) {
String answer = chatClient.prompt()
.advisors(new QuestionAnswerAdvisor(vectorStore))
.user(question).call().content();
return answer;
}
}
loadDocs
メソッドでAzure Cosmos DBへドキュメントを格納しています。ベクトル化は内部で自動的に行われます。
rag
メソッドでRAGを利用して会話します。
動作確認
アプリケーションを起動します。
mvn spring-boot:run
まずはAzure Cosmos DBへドキュメントを格納していない状態で質問してみましょう。
curl localhost:8080/rag -s -d question="where did harrison work?"
すると回答できない旨のメッセージが返されます。
I'm sorry, but the context does not provide information about where Harrison worked.
それではAzure Cosmos DBへドキュメントを格納してみましょう。
curl localhost:8080/load-docs -X POST
再度、質問します。
curl localhost:8080/rag -s -d question="where did harrison work?"
今度は回答が返されました。
Harrison worked at Kensho.
Azureリソースの後始末
使い終わったら後始末しましょう。
リソースグループを削除します。
az group delete --name $resource_group
終わりに
Azure Cosmos DBにベクトル検索の機能がサポートされたことと、それをSpring AIから簡単に使えることを確認しました。
さらにAzure Cosmos DBは現時点ではプレビューですが全文検索やハイブリッド検索もサポートするようになり、Azure AI Searchの役割を置き換えられそうになってきました。
管理するリソースは少ないに越したことはないので、すべてAzure Cosmos DBで済ませられるなら嬉しいですね。
Discussion