Zenn
🌿

Spring AIでAzure Cosmos DBのベクトル検索を試す

2024/12/25に公開

概要

Azure Cosmos DBがベクトル検索できるようになり、Spring AIとも統合されました。

https://devblogs.microsoft.com/cosmosdb/announcing-azure-cosmos-db-integration-with-spring-ai-and-langchain4j/

簡単なRAGを通じてSpring AIとAzure Cosmos DB、Azure OpenAI Serviceを用いたベクトル検索を試してみます。

コード例

今回のコード例は次のリポジトリのspring-ai-cosmos-vector-store-exampleというディレクトリにあります。advent-calendar-20241225というタグを参照してください。

https://github.com/backpaper0/spring-ai-example/tree/advent-calendar-20241225/spring-ai-cosmos-vector-store-example

以降の説明はこのコード例をもとにして記載しています。

バージョンなど

  • 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を利用

データベースやコンテナーは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のコードを流用します。

https://zenn.dev/backpaper0/articles/05f4b62d5bed77

コントローラーのコードを掲載しておきます。

@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の役割を置き換えられそうになってきました。

https://learn.microsoft.com/ja-jp/azure/cosmos-db/gen-ai/full-text-search

https://learn.microsoft.com/ja-jp/azure/cosmos-db/gen-ai/hybrid-search

管理するリソースは少ないに越したことはないので、すべてAzure Cosmos DBで済ませられるなら嬉しいですね。

Discussion

ログインするとコメントできます