Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Make chat history available to RetrievalAugmentationAdvisor

Browse Source 
* Extend Query with conversation history in RetrievalAugmentationAdvisor
* Add integration tests for query compression and rewrite

Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>
This commit is contained in:
Thomas Vitale
2024-12-24 15:06:24 +01:00
committed by Mark Pollack
parent 387de95ac5
commit 4a90a46925
2 changed files with 80 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisor.java
View File

@@ -102,8 +102,11 @@ public final class RetrievalAugmentationAdvisor implements BaseAdvisor {
public AdvisedRequest before(AdvisedRequest request) {
Map<String, Object> context = new HashMap<>(request.adviseContext());
// 0. Create a query from the user text and parameters.
Query originalQuery = new Query(new PromptTemplate(request.userText(), request.userParams()).render());
// 0. Create a query from the user text, parameters, and conversation history.
Query originalQuery = Query.builder()
.text(new PromptTemplate(request.userText(), request.userParams()).render())
.history(request.messages())
.build();
// 1. Transform original user query based on a chain of query transformers.
Query transformedQuery = originalQuery;

75
spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java
View File

@@ -24,7 +24,10 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
@@ -34,6 +37,8 @@ import org.springframework.ai.evaluation.RelevancyEvaluator;
import org.springframework.ai.integration.tests.TestApplication;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.CompressionQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
@@ -103,6 +108,76 @@ class RetrievalAugmentationAdvisorIT {
evaluateRelevancy(question, chatResponse);
}
@Test
void ragWithCompression() {
MessageChatMemoryAdvisor memoryAdvisor = MessageChatMemoryAdvisor.builder(new InMemoryChatMemory()).build();
RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder()
.queryTransformers(CompressionQueryTransformer.builder()
.chatClientBuilder(ChatClient.builder(this.openAiChatModel))
.build())
.documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(this.pgVectorStore).build())
.build();
ChatClient chatClient = ChatClient.builder(this.openAiChatModel)
.defaultAdvisors(memoryAdvisor, ragAdvisor)
.build();
String conversationId = "007";
ChatResponse chatResponse1 = chatClient.prompt()
.user("Where does the adventure of Anacletus and Birba take place?")
.advisors(advisors -> advisors.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,
conversationId))
.call()
.chatResponse();
assertThat(chatResponse1).isNotNull();
String response1 = chatResponse1.getResult().getOutput().getText();
System.out.println(response1);
ChatResponse chatResponse2 = chatClient.prompt()
.user("Did they meet any cow?")
.advisors(advisors -> advisors.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,
conversationId))
.call()
.chatResponse();
assertThat(chatResponse2).isNotNull();
String response2 = chatResponse2.getResult().getOutput().getText();
System.out.println(response2);
assertThat(response2.toLowerCase()).containsIgnoringCase("Fergus");
}
@Test
void ragWithRewrite() {
String question = "Where are the main characters going?";
RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder()
.queryTransformers(RewriteQueryTransformer.builder()
.chatClientBuilder(ChatClient.builder(this.openAiChatModel))
.targetSearchSystem("vector store")
.build())
.documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(this.pgVectorStore).build())
.build();
ChatResponse chatResponse = ChatClient.builder(this.openAiChatModel)
.build()
.prompt()
.user(question)
.advisors(ragAdvisor)
.call()
.chatResponse();
assertThat(chatResponse).isNotNull();
String response = chatResponse.getResult().getOutput().getText();
System.out.println(response);
assertThat(response).containsIgnoringCase("Loch of the Stars");
evaluateRelevancy(question, chatResponse);
}
@Test
void ragWithTranslation() {
String question = "Hvor finder Anacletus og Birbas eventyr sted?";