🌿

Spring AIでRAGするのが前より簡単になってた

2024/12/01に公開

概要

今年の春ごろに「Spring AIでLangChainのクックブックにあるRAGをなぞる」という記事を書きました。

https://zenn.dev/backpaper0/articles/ee37fd39c8feff

あれからSpring AIも進化し、より簡単にRAGを行えるようになったので試してみます。

バージョンなど

  • Java 21
  • Spring Boot 3.4.0
  • Spring AI 1.0.0-M4

簡単にRAGを行う

まずコードを示します。
Advisorという仕組みが入ったことによって、とても簡単なコードでRAGを行えるようになりました。
QuestionAnswerAdvisorがベクトルデータベースへの検索とプロンプト構築を行なってくれます。

@PostMapping("/rag1")
public Object rag1(@RequestParam String question) {

    String answer = chatClient.prompt()
            .advisors(new QuestionAnswerAdvisor(vectorStore))
            .user(question).call().content();

    return answer;
}

ベクトル検索で取得した情報を利用するプロンプトは暗黙的にデフォルトのものが設定されていますが、QuestionAnswerAdvisorのコンストラクタ引数で変更できます。

比較のため冒頭に紹介した過去の記事で解説しているコードも示します。
検索と生成を素朴に書き下していることがわかります。

@PostMapping("/rag1")
public Object rag1(@RequestParam String question) {

    // 検索
    List<Document> docs = vectorStore.similaritySearch(question);
    Iterator<Document> iter = docs.iterator();
    if (!iter.hasNext()) {
        // 検索に何もヒットしなかった
        return "I do not know.";
    }

    // 生成
    String context = iter.next().getContent();
    String prompt = """
            Answer the question based only on the following context:
            %2$s

            Question: %1$s
            """.formatted(question, context);
    String answer = chatClient.prompt().user(prompt).call().content();

    return answer;
}

履歴を伴う会話も簡単になった

RAGのことだけさらっと書いて終わろうと思っていたのですが、Advisorの実装クラスを見ていると履歴を伴う会話も簡単になっていることに感動しました。

次に示すコードで履歴を伴う会話ができます。
MessageChatMemoryAdvisorが履歴を管理してくれます。履歴のキーがconversationIdです。

@PostMapping("/chat")
public Object chat(@RequestParam String query, @RequestParam String conversationId) {
    String answer = chatClient.prompt()
            .advisors(new MessageChatMemoryAdvisor(chatMemory))
            .advisors(advisor -> advisor
                    .param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId))
            .user(query)
            .call().content();
    return answer;
}

会話を試してみました。
履歴を踏まえて会話が成り立っていることや、履歴のキー(conversationId)が変わると記憶喪失(?)になっていることがわかります。

$ curl localhost:8080/chat -d conversationId=one -d query="2に3を掛けて"
2に3を掛けると、6になります。
$ curl localhost:8080/chat -d conversationId=one -d query="4を足して"
6に4を足すと、10になります。
$ curl localhost:8080/chat -d conversationId=two -d query="4を足して"
「4を足して」とは、何に4を足したいのか具体的に教えていただけますか?例えば、特定の数字に4を足す場合や、計算の問題について教えてください。

なお、履歴のリポジトリであるChatMemoryを別途コンポーネント定義する必要があります。

@Bean
ChatMemory chatMemory() {
    return new InMemoryChatMemory();
}

実装は今のところインメモリとCassandraの2つらしいです。(JDBC実装が作られる機運の高まりを感じる……!

終わりに

Advisorという仕組みでRAGや履歴を伴う会話が簡単に実装できることがわかりました。
また、これらを組み合わせることで履歴とRAGを伴う会話なんかも簡単に実装できそうです。

LangChainはRunnableが最小単位の処理になっておりRunnableをチェーンさせることでロジックを構築しますが、Spring AIではAdvisorを重ねることでロジックを構築することになりそうですね。
Advisorの考え方は従来のSpring Frameworkにもあるもので馴染み深く、小さな処理を組み合わせる方法として相性が良いと感じました。

久しぶりにSpring AIで遊びましたが、楽しかった〜!

今回のコードはGitHubにも置いています。advent-calendar-20241201というタグを打っています。

https://github.com/backpaper0/spring-ai-example/releases/tag/advent-calendar-20241201

参考

Discussion