diff --git a/langchain4j-spring-boot-starter/src/test/java/dev/langchain4j/service/spring/mode/automatic/withRetrievalAugmentor/AiServiceWithRetrievalAugmentorApplication.java b/langchain4j-spring-boot-starter/src/test/java/dev/langchain4j/service/spring/mode/automatic/withRetrievalAugmentor/AiServiceWithRetrievalAugmentorApplication.java
index 6721d41..38437bc 100644
--- a/langchain4j-spring-boot-starter/src/test/java/dev/langchain4j/service/spring/mode/automatic/withRetrievalAugmentor/AiServiceWithRetrievalAugmentorApplication.java
+++ b/langchain4j-spring-boot-starter/src/test/java/dev/langchain4j/service/spring/mode/automatic/withRetrievalAugmentor/AiServiceWithRetrievalAugmentorApplication.java
@@ -1,6 +1,7 @@
 package dev.langchain4j.service.spring.mode.automatic.withRetrievalAugmentor;
 
 import dev.langchain4j.data.message.UserMessage;
+import dev.langchain4j.rag.AugmentationResult;
 import dev.langchain4j.rag.RetrievalAugmentor;
 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
@@ -11,7 +12,9 @@ class AiServiceWithRetrievalAugmentorApplication {
 
     @Bean
     RetrievalAugmentor retrievalAugmentor() {
-        return (userMessage, metadata) -> UserMessage.from("My name is Klaus." + userMessage);
+        return augmentationRequest -> AugmentationResult.builder()
+                .chatMessage(UserMessage.from("My name is Klaus." + augmentationRequest.chatMessage()))
+                .build();
     }
 
     public static void main(String[] args) {