|
24 | 24 | import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
25 | 25 |
|
26 | 26 | import org.springframework.ai.chat.client.ChatClient;
|
| 27 | +import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor; |
| 28 | +import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor; |
27 | 29 | import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
|
| 30 | +import org.springframework.ai.chat.memory.InMemoryChatMemory; |
28 | 31 | import org.springframework.ai.chat.model.ChatResponse;
|
29 | 32 | import org.springframework.ai.document.Document;
|
30 | 33 | import org.springframework.ai.document.DocumentReader;
|
|
34 | 37 | import org.springframework.ai.integration.tests.TestApplication;
|
35 | 38 | import org.springframework.ai.openai.OpenAiChatModel;
|
36 | 39 | import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
|
| 40 | +import org.springframework.ai.rag.preretrieval.query.transformation.CompressionQueryTransformer; |
| 41 | +import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer; |
37 | 42 | import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
|
38 | 43 | import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
|
39 | 44 | import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
|
@@ -103,6 +108,74 @@ void ragBasic() {
|
103 | 108 | evaluateRelevancy(question, chatResponse);
|
104 | 109 | }
|
105 | 110 |
|
| 111 | + @Test |
| 112 | + void ragWithCompression() { |
| 113 | + MessageChatMemoryAdvisor memoryAdvisor = MessageChatMemoryAdvisor.builder(new InMemoryChatMemory()).build(); |
| 114 | + |
| 115 | + RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder() |
| 116 | + .queryTransformers(CompressionQueryTransformer.builder() |
| 117 | + .chatClientBuilder(ChatClient.builder(this.openAiChatModel)) |
| 118 | + .build()) |
| 119 | + .documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(this.pgVectorStore).build()) |
| 120 | + .build(); |
| 121 | + |
| 122 | + ChatClient chatClient = ChatClient.builder(this.openAiChatModel) |
| 123 | + .defaultAdvisors(memoryAdvisor, ragAdvisor) |
| 124 | + .build(); |
| 125 | + |
| 126 | + String conversationId = "007"; |
| 127 | + |
| 128 | + ChatResponse chatResponse1 = chatClient.prompt() |
| 129 | + .user("Where does the adventure of Anacletus and Birba take place?") |
| 130 | + .advisors(advisors -> advisors.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, |
| 131 | + conversationId)) |
| 132 | + .call() |
| 133 | + .chatResponse(); |
| 134 | + |
| 135 | + assertThat(chatResponse1).isNotNull(); |
| 136 | + String response1 = chatResponse1.getResult().getOutput().getText(); |
| 137 | + System.out.println(response1); |
| 138 | + |
| 139 | + ChatResponse chatResponse2 = chatClient.prompt() |
| 140 | + .user("Did they meet any cow?") |
| 141 | + .advisors(advisors -> advisors.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, |
| 142 | + conversationId)) |
| 143 | + .call() |
| 144 | + .chatResponse(); |
| 145 | + |
| 146 | + assertThat(chatResponse2).isNotNull(); |
| 147 | + String response2 = chatResponse2.getResult().getOutput().getText(); |
| 148 | + System.out.println(response2); |
| 149 | + assertThat(response2.toLowerCase()).containsIgnoringCase("Fergus"); |
| 150 | + } |
| 151 | + |
| 152 | + @Test |
| 153 | + void ragWithRewrite() { |
| 154 | + String question = "Where are the main characters going?"; |
| 155 | + |
| 156 | + RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder() |
| 157 | + .queryTransformers(RewriteQueryTransformer.builder() |
| 158 | + .chatClientBuilder(ChatClient.builder(this.openAiChatModel)) |
| 159 | + .targetSearchSystem("vector store") |
| 160 | + .build()) |
| 161 | + .documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(this.pgVectorStore).build()) |
| 162 | + .build(); |
| 163 | + |
| 164 | + ChatResponse chatResponse = ChatClient.builder(this.openAiChatModel).build().prompt() |
| 165 | + .user(question) |
| 166 | + .advisors(ragAdvisor) |
| 167 | + .call() |
| 168 | + .chatResponse(); |
| 169 | + |
| 170 | + assertThat(chatResponse).isNotNull(); |
| 171 | + |
| 172 | + String response = chatResponse.getResult().getOutput().getText(); |
| 173 | + System.out.println(response); |
| 174 | + assertThat(response).containsIgnoringCase("Loch of the Stars"); |
| 175 | + |
| 176 | + evaluateRelevancy(question, chatResponse); |
| 177 | + } |
| 178 | + |
106 | 179 | @Test
|
107 | 180 | void ragWithTranslation() {
|
108 | 181 | String question = "Hvor finder Anacletus og Birbas eventyr sted?";
|
|
0 commit comments