refactor: Extract vector store and advisor functionality from spring-ai-core
Major Changes: - Created new module spring-ai-vector-store from spring-ai-core functionality - Split advisor functionality into three new modules: * advisor-memory: Memory-based chat advisors * advisor-rag: Retrieval Augmentation Generation advisors * advisor-vector-store: Vector store based advisors
This commit is contained in:
Mark Pollack
@@ -0,0 +1,356 @@
|
||||
/*
|
||||
* Copyright 2023-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.chat.client.advisor.vectorstore;
|
||||
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.function.Predicate;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
|
||||
import org.springframework.ai.chat.client.advisor.api.AdvisedRequest;
|
||||
import org.springframework.ai.chat.client.advisor.api.AdvisedResponse;
|
||||
import org.springframework.ai.chat.client.advisor.api.CallAroundAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.CallAroundAdvisorChain;
|
||||
import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisorChain;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.prompt.PromptTemplate;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser;
|
||||
import org.springframework.util.Assert;
|
||||
import org.springframework.util.StringUtils;
|
||||
|
||||
/**
|
||||
* Context for the question is retrieved from a Vector Store and added to the prompt's
|
||||
* user text.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Timo Salm
|
||||
* @author Ilayaperumal Gopinathan
|
||||
* @since 1.0.0
|
||||
*/
|
||||
public class QuestionAnswerAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {
|
||||
|
||||
public static final String RETRIEVED_DOCUMENTS = "qa_retrieved_documents";
|
||||
|
||||
public static final String FILTER_EXPRESSION = "qa_filter_expression";
|
||||
|
||||
private static final String DEFAULT_USER_TEXT_ADVISE = """
|
||||
|
||||
Context information is below, surrounded by ---------------------
|
||||
|
||||
---------------------
|
||||
{question_answer_context}
|
||||
---------------------
|
||||
|
||||
Given the context and provided history information and not prior knowledge,
|
||||
reply to the user comment. If the answer is not in the context, inform
|
||||
the user that you can't answer the question.
|
||||
""";
|
||||
|
||||
private static final int DEFAULT_ORDER = 0;
|
||||
|
||||
private final VectorStore vectorStore;
|
||||
|
||||
private final String userTextAdvise;
|
||||
|
||||
private final SearchRequest searchRequest;
|
||||
|
||||
private final boolean protectFromBlocking;
|
||||
|
||||
private final int order;
|
||||
|
||||
/**
|
||||
* The QuestionAnswerAdvisor retrieves context information from a Vector Store and
|
||||
* combines it with the user's text.
|
||||
* @param vectorStore The vector store to use
|
||||
*/
|
||||
public QuestionAnswerAdvisor(VectorStore vectorStore) {
|
||||
this(vectorStore, SearchRequest.defaults(), DEFAULT_USER_TEXT_ADVISE);
|
||||
}
|
||||
|
||||
/**
|
||||
* The QuestionAnswerAdvisor retrieves context information from a Vector Store and
|
||||
* combines it with the user's text.
|
||||
* @param vectorStore The vector store to use
|
||||
* @param searchRequest The search request defined using the portable filter
|
||||
* expression syntax
|
||||
*/
|
||||
public QuestionAnswerAdvisor(VectorStore vectorStore, SearchRequest searchRequest) {
|
||||
this(vectorStore, searchRequest, DEFAULT_USER_TEXT_ADVISE);
|
||||
}
|
||||
|
||||
/**
|
||||
* The QuestionAnswerAdvisor retrieves context information from a Vector Store and
|
||||
* combines it with the user's text.
|
||||
* @param vectorStore The vector store to use
|
||||
* @param searchRequest The search request defined using the portable filter
|
||||
* expression syntax
|
||||
* @param userTextAdvise The user text to append to the existing user prompt. The text
|
||||
* should contain a placeholder named "question_answer_context".
|
||||
*/
|
||||
public QuestionAnswerAdvisor(VectorStore vectorStore, SearchRequest searchRequest, String userTextAdvise) {
|
||||
this(vectorStore, searchRequest, userTextAdvise, true);
|
||||
}
|
||||
|
||||
/**
|
||||
* The QuestionAnswerAdvisor retrieves context information from a Vector Store and
|
||||
* combines it with the user's text.
|
||||
* @param vectorStore The vector store to use
|
||||
* @param searchRequest The search request defined using the portable filter
|
||||
* expression syntax
|
||||
* @param userTextAdvise The user text to append to the existing user prompt. The text
|
||||
* should contain a placeholder named "question_answer_context".
|
||||
* @param protectFromBlocking If true the advisor will protect the execution from
|
||||
* blocking threads. If false the advisor will not protect the execution from blocking
|
||||
* threads. This is useful when the advisor is used in a non-blocking environment. It
|
||||
* is true by default.
|
||||
*/
|
||||
public QuestionAnswerAdvisor(VectorStore vectorStore, SearchRequest searchRequest, String userTextAdvise,
|
||||
boolean protectFromBlocking) {
|
||||
this(vectorStore, searchRequest, userTextAdvise, protectFromBlocking, DEFAULT_ORDER);
|
||||
}
|
||||
|
||||
/**
|
||||
* The QuestionAnswerAdvisor retrieves context information from a Vector Store and
|
||||
* combines it with the user's text.
|
||||
* @param vectorStore The vector store to use
|
||||
* @param searchRequest The search request defined using the portable filter
|
||||
* expression syntax
|
||||
* @param userTextAdvise The user text to append to the existing user prompt. The text
|
||||
* should contain a placeholder named "question_answer_context".
|
||||
* @param protectFromBlocking If true the advisor will protect the execution from
|
||||
* blocking threads. If false the advisor will not protect the execution from blocking
|
||||
* threads. This is useful when the advisor is used in a non-blocking environment. It
|
||||
* is true by default.
|
||||
* @param order The order of the advisor.
|
||||
*/
|
||||
public QuestionAnswerAdvisor(VectorStore vectorStore, SearchRequest searchRequest, String userTextAdvise,
|
||||
boolean protectFromBlocking, int order) {
|
||||
|
||||
Assert.notNull(vectorStore, "The vectorStore must not be null!");
|
||||
Assert.notNull(searchRequest, "The searchRequest must not be null!");
|
||||
Assert.hasText(userTextAdvise, "The userTextAdvise must not be empty!");
|
||||
|
||||
this.vectorStore = vectorStore;
|
||||
this.searchRequest = searchRequest;
|
||||
this.userTextAdvise = userTextAdvise;
|
||||
this.protectFromBlocking = protectFromBlocking;
|
||||
this.order = order;
|
||||
}
|
||||
|
||||
public static Builder builder(VectorStore vectorStore) {
|
||||
return new Builder(vectorStore);
|
||||
}
|
||||
|
||||
@Override
|
||||
public String getName() {
|
||||
return this.getClass().getSimpleName();
|
||||
}
|
||||
|
||||
@Override
|
||||
public int getOrder() {
|
||||
return this.order;
|
||||
}
|
||||
|
||||
@Override
|
||||
public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
|
||||
|
||||
AdvisedRequest advisedRequest2 = before(advisedRequest);
|
||||
|
||||
AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest2);
|
||||
|
||||
return after(advisedResponse);
|
||||
}
|
||||
|
||||
@Override
|
||||
public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
|
||||
|
||||
// This can be executed by both blocking and non-blocking Threads
|
||||
// E.g. a command line or Tomcat blocking Thread implementation
|
||||
// or by a WebFlux dispatch in a non-blocking manner.
|
||||
Flux<AdvisedResponse> advisedResponses = (this.protectFromBlocking) ?
|
||||
// @formatter:off
|
||||
Mono.just(advisedRequest)
|
||||
.publishOn(Schedulers.boundedElastic())
|
||||
.map(this::before)
|
||||
.flatMapMany(request -> chain.nextAroundStream(request))
|
||||
: chain.nextAroundStream(before(advisedRequest));
|
||||
// @formatter:on
|
||||
|
||||
return advisedResponses.map(ar -> {
|
||||
if (onFinishReason().test(ar)) {
|
||||
ar = after(ar);
|
||||
}
|
||||
return ar;
|
||||
});
|
||||
}
|
||||
|
||||
private AdvisedRequest before(AdvisedRequest request) {
|
||||
|
||||
var context = new HashMap<>(request.adviseContext());
|
||||
|
||||
// 1. Advise the system text.
|
||||
String advisedUserText = request.userText() + System.lineSeparator() + this.userTextAdvise;
|
||||
|
||||
// 2. Search for similar documents in the vector store.
|
||||
String query = new PromptTemplate(request.userText(), request.userParams()).render();
|
||||
var searchRequestToUse = SearchRequest.from(this.searchRequest)
|
||||
.withQuery(query)
|
||||
.withFilterExpression(doGetFilterExpression(context));
|
||||
|
||||
List<Document> documents = this.vectorStore.similaritySearch(searchRequestToUse);
|
||||
|
||||
// 3. Create the context from the documents.
|
||||
context.put(RETRIEVED_DOCUMENTS, documents);
|
||||
|
||||
String documentContext = documents.stream()
|
||||
.map(Document::getText)
|
||||
.collect(Collectors.joining(System.lineSeparator()));
|
||||
|
||||
// 4. Advise the user parameters.
|
||||
Map<String, Object> advisedUserParams = new HashMap<>(request.userParams());
|
||||
advisedUserParams.put("question_answer_context", documentContext);
|
||||
|
||||
AdvisedRequest advisedRequest = AdvisedRequest.from(request)
|
||||
.userText(advisedUserText)
|
||||
.userParams(advisedUserParams)
|
||||
.adviseContext(context)
|
||||
.build();
|
||||
|
||||
return advisedRequest;
|
||||
}
|
||||
|
||||
private AdvisedResponse after(AdvisedResponse advisedResponse) {
|
||||
ChatResponse.Builder chatResponseBuilder = ChatResponse.builder().from(advisedResponse.response());
|
||||
chatResponseBuilder.withMetadata(RETRIEVED_DOCUMENTS, advisedResponse.adviseContext().get(RETRIEVED_DOCUMENTS));
|
||||
return new AdvisedResponse(chatResponseBuilder.build(), advisedResponse.adviseContext());
|
||||
}
|
||||
|
||||
protected Filter.Expression doGetFilterExpression(Map<String, Object> context) {
|
||||
|
||||
if (!context.containsKey(FILTER_EXPRESSION)
|
||||
|| !StringUtils.hasText(context.get(FILTER_EXPRESSION).toString())) {
|
||||
return this.searchRequest.getFilterExpression();
|
||||
}
|
||||
return new FilterExpressionTextParser().parse(context.get(FILTER_EXPRESSION).toString());
|
||||
|
||||
}
|
||||
|
||||
private Predicate<AdvisedResponse> onFinishReason() {
|
||||
return advisedResponse -> advisedResponse.response()
|
||||
.getResults()
|
||||
.stream()
|
||||
.filter(result -> result != null && result.getMetadata() != null
|
||||
&& StringUtils.hasText(result.getMetadata().getFinishReason()))
|
||||
.findFirst()
|
||||
.isPresent();
|
||||
}
|
||||
|
||||
public static final class Builder {
|
||||
|
||||
private final VectorStore vectorStore;
|
||||
|
||||
private SearchRequest searchRequest = SearchRequest.defaults();
|
||||
|
||||
private String userTextAdvise = DEFAULT_USER_TEXT_ADVISE;
|
||||
|
||||
private boolean protectFromBlocking = true;
|
||||
|
||||
private int order = DEFAULT_ORDER;
|
||||
|
||||
private Builder(VectorStore vectorStore) {
|
||||
Assert.notNull(vectorStore, "The vectorStore must not be null!");
|
||||
this.vectorStore = vectorStore;
|
||||
}
|
||||
|
||||
public Builder searchRequest(SearchRequest searchRequest) {
|
||||
Assert.notNull(searchRequest, "The searchRequest must not be null!");
|
||||
this.searchRequest = searchRequest;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder userTextAdvise(String userTextAdvise) {
|
||||
Assert.hasText(userTextAdvise, "The userTextAdvise must not be empty!");
|
||||
this.userTextAdvise = userTextAdvise;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder protectFromBlocking(boolean protectFromBlocking) {
|
||||
this.protectFromBlocking = protectFromBlocking;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder order(int order) {
|
||||
this.order = order;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated use {@link #searchRequest(SearchRequest)} instead.
|
||||
*/
|
||||
@Deprecated(forRemoval = true, since = "1.0.0-M5")
|
||||
public Builder withSearchRequest(SearchRequest searchRequest) {
|
||||
Assert.notNull(searchRequest, "The searchRequest must not be null!");
|
||||
this.searchRequest = searchRequest;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated use {@link #userTextAdvise(String)} instead.
|
||||
*/
|
||||
@Deprecated(forRemoval = true, since = "1.0.0-M5")
|
||||
public Builder withUserTextAdvise(String userTextAdvise) {
|
||||
Assert.hasText(userTextAdvise, "The userTextAdvise must not be empty!");
|
||||
this.userTextAdvise = userTextAdvise;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated use {@link #protectFromBlocking(boolean)} instead.
|
||||
*/
|
||||
@Deprecated(forRemoval = true, since = "1.0.0-M5")
|
||||
public Builder withProtectFromBlocking(boolean protectFromBlocking) {
|
||||
this.protectFromBlocking = protectFromBlocking;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated use {@link #order(int)} instead.
|
||||
*/
|
||||
@Deprecated(forRemoval = true, since = "1.0.0-M5")
|
||||
public Builder withOrder(int order) {
|
||||
this.order = order;
|
||||
return this;
|
||||
}
|
||||
|
||||
public QuestionAnswerAdvisor build() {
|
||||
return new QuestionAnswerAdvisor(this.vectorStore, this.searchRequest, this.userTextAdvise,
|
||||
this.protectFromBlocking, this.order);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
/*
|
||||
* Copyright 2023-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.chat.client.advisor.vectorstore;
|
||||
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import reactor.core.publisher.Flux;
|
||||
|
||||
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.AdvisedRequest;
|
||||
import org.springframework.ai.chat.client.advisor.api.AdvisedResponse;
|
||||
import org.springframework.ai.chat.client.advisor.api.Advisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.CallAroundAdvisorChain;
|
||||
import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisorChain;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.messages.Message;
|
||||
import org.springframework.ai.chat.messages.MessageType;
|
||||
import org.springframework.ai.chat.messages.UserMessage;
|
||||
import org.springframework.ai.chat.model.MessageAggregator;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
import org.springframework.util.StringUtils;
|
||||
|
||||
/**
|
||||
* Memory is retrieved from a VectorStore added into the prompt's system text.
|
||||
*
|
||||
* This only works for text based exchanges with the models, not multi-modal exchanges.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Thomas Vitale
|
||||
* @since 1.0.0
|
||||
*/
|
||||
public class VectorStoreChatMemoryAdvisor extends AbstractChatMemoryAdvisor<VectorStore> {
|
||||
|
||||
private static final String DOCUMENT_METADATA_CONVERSATION_ID = "conversationId";
|
||||
|
||||
private static final String DOCUMENT_METADATA_MESSAGE_TYPE = "messageType";
|
||||
|
||||
private static final String DEFAULT_SYSTEM_TEXT_ADVISE = """
|
||||
|
||||
Use the long term conversation memory from the LONG_TERM_MEMORY section to provide accurate answers.
|
||||
|
||||
---------------------
|
||||
LONG_TERM_MEMORY:
|
||||
{long_term_memory}
|
||||
---------------------
|
||||
""";
|
||||
|
||||
private final String systemTextAdvise;
|
||||
|
||||
public VectorStoreChatMemoryAdvisor(VectorStore vectorStore) {
|
||||
this(vectorStore, DEFAULT_SYSTEM_TEXT_ADVISE);
|
||||
}
|
||||
|
||||
public VectorStoreChatMemoryAdvisor(VectorStore vectorStore, String systemTextAdvise) {
|
||||
super(vectorStore);
|
||||
this.systemTextAdvise = systemTextAdvise;
|
||||
}
|
||||
|
||||
public VectorStoreChatMemoryAdvisor(VectorStore vectorStore, String defaultConversationId,
|
||||
int chatHistoryWindowSize) {
|
||||
this(vectorStore, defaultConversationId, chatHistoryWindowSize, DEFAULT_SYSTEM_TEXT_ADVISE);
|
||||
}
|
||||
|
||||
public VectorStoreChatMemoryAdvisor(VectorStore vectorStore, String defaultConversationId,
|
||||
int chatHistoryWindowSize, String systemTextAdvise) {
|
||||
this(vectorStore, defaultConversationId, chatHistoryWindowSize, systemTextAdvise,
|
||||
Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER);
|
||||
}
|
||||
|
||||
/**
|
||||
* Constructor for VectorStoreChatMemoryAdvisor.
|
||||
* @param vectorStore the vector store instance used for managing and querying
|
||||
* documents.
|
||||
* @param defaultConversationId the default conversation ID used if none is provided
|
||||
* in the context.
|
||||
* @param chatHistoryWindowSize the window size for the chat history retrieval.
|
||||
* @param systemTextAdvise the system text advice used for the chat advisor system.
|
||||
* @param order the order of precedence for this advisor in the chain.
|
||||
*/
|
||||
public VectorStoreChatMemoryAdvisor(VectorStore vectorStore, String defaultConversationId,
|
||||
int chatHistoryWindowSize, String systemTextAdvise, int order) {
|
||||
super(vectorStore, defaultConversationId, chatHistoryWindowSize, true, order);
|
||||
this.systemTextAdvise = systemTextAdvise;
|
||||
}
|
||||
|
||||
public static Builder builder(VectorStore chatMemory) {
|
||||
return new Builder(chatMemory);
|
||||
}
|
||||
|
||||
@Override
|
||||
public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
|
||||
|
||||
advisedRequest = this.before(advisedRequest);
|
||||
|
||||
AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);
|
||||
|
||||
this.observeAfter(advisedResponse);
|
||||
|
||||
return advisedResponse;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
|
||||
|
||||
Flux<AdvisedResponse> advisedResponses = this.doNextWithProtectFromBlockingBefore(advisedRequest, chain,
|
||||
this::before);
|
||||
|
||||
// The observeAfter will certainly be executed on non-blocking Threads in case
|
||||
// of some models - e.g. when the model client is a WebClient
|
||||
return new MessageAggregator().aggregateAdvisedResponse(advisedResponses, this::observeAfter);
|
||||
}
|
||||
|
||||
private AdvisedRequest before(AdvisedRequest request) {
|
||||
|
||||
String advisedSystemText;
|
||||
if (StringUtils.hasText(request.systemText())) {
|
||||
advisedSystemText = request.systemText() + System.lineSeparator() + this.systemTextAdvise;
|
||||
}
|
||||
else {
|
||||
advisedSystemText = this.systemTextAdvise;
|
||||
}
|
||||
|
||||
var searchRequest = SearchRequest.query(request.userText())
|
||||
.withTopK(this.doGetChatMemoryRetrieveSize(request.adviseContext()))
|
||||
.withFilterExpression(DOCUMENT_METADATA_CONVERSATION_ID + "=='"
|
||||
+ this.doGetConversationId(request.adviseContext()) + "'");
|
||||
|
||||
List<Document> documents = this.getChatMemoryStore().similaritySearch(searchRequest);
|
||||
|
||||
String longTermMemory = documents.stream()
|
||||
.map(Document::getText)
|
||||
.collect(Collectors.joining(System.lineSeparator()));
|
||||
|
||||
Map<String, Object> advisedSystemParams = new HashMap<>(request.systemParams());
|
||||
advisedSystemParams.put("long_term_memory", longTermMemory);
|
||||
|
||||
AdvisedRequest advisedRequest = AdvisedRequest.from(request)
|
||||
.systemText(advisedSystemText)
|
||||
.systemParams(advisedSystemParams)
|
||||
.build();
|
||||
|
||||
UserMessage userMessage = new UserMessage(request.userText(), request.media());
|
||||
this.getChatMemoryStore()
|
||||
.write(toDocuments(List.of(userMessage), this.doGetConversationId(request.adviseContext())));
|
||||
|
||||
return advisedRequest;
|
||||
}
|
||||
|
||||
private void observeAfter(AdvisedResponse advisedResponse) {
|
||||
|
||||
List<Message> assistantMessages = advisedResponse.response()
|
||||
.getResults()
|
||||
.stream()
|
||||
.map(g -> (Message) g.getOutput())
|
||||
.toList();
|
||||
|
||||
this.getChatMemoryStore()
|
||||
.write(toDocuments(assistantMessages, this.doGetConversationId(advisedResponse.adviseContext())));
|
||||
}
|
||||
|
||||
private List<Document> toDocuments(List<Message> messages, String conversationId) {
|
||||
|
||||
List<Document> docs = messages.stream()
|
||||
.filter(m -> m.getMessageType() == MessageType.USER || m.getMessageType() == MessageType.ASSISTANT)
|
||||
.map(message -> {
|
||||
var metadata = new HashMap<>(message.getMetadata() != null ? message.getMetadata() : new HashMap<>());
|
||||
metadata.put(DOCUMENT_METADATA_CONVERSATION_ID, conversationId);
|
||||
metadata.put(DOCUMENT_METADATA_MESSAGE_TYPE, message.getMessageType().name());
|
||||
if (message instanceof UserMessage userMessage) {
|
||||
return Document.builder()
|
||||
.text(userMessage.getText())
|
||||
// userMessage.getMedia().get(0).getId()
|
||||
// TODO vector store for memory would not store this into the
|
||||
// vector store, could store an 'id' instead
|
||||
// .media(userMessage.getMedia())
|
||||
.metadata(metadata)
|
||||
.build();
|
||||
}
|
||||
else if (message instanceof AssistantMessage assistantMessage) {
|
||||
return Document.builder().text(assistantMessage.getText()).metadata(metadata).build();
|
||||
}
|
||||
throw new RuntimeException("Unknown message type: " + message.getMessageType());
|
||||
})
|
||||
.toList();
|
||||
|
||||
return docs;
|
||||
}
|
||||
|
||||
public static class Builder extends AbstractChatMemoryAdvisor.AbstractBuilder<VectorStore> {
|
||||
|
||||
private String systemTextAdvise = DEFAULT_SYSTEM_TEXT_ADVISE;
|
||||
|
||||
protected Builder(VectorStore chatMemory) {
|
||||
super(chatMemory);
|
||||
}
|
||||
|
||||
public Builder systemTextAdvise(String systemTextAdvise) {
|
||||
this.systemTextAdvise = systemTextAdvise;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated use {@link #systemTextAdvise(String)} instead.
|
||||
*/
|
||||
@Deprecated(forRemoval = true, since = "1.0.0-M5")
|
||||
public Builder withSystemTextAdvise(String systemTextAdvise) {
|
||||
this.systemTextAdvise = systemTextAdvise;
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public VectorStoreChatMemoryAdvisor build() {
|
||||
return new VectorStoreChatMemoryAdvisor(this.chatMemory, this.conversationId, this.chatMemoryRetrieveSize,
|
||||
this.systemTextAdvise);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
/*
|
||||
* Copyright 2023-2025 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.chat.client.advisor.vectorstore;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.ExtendWith;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.Captor;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.junit.jupiter.MockitoExtension;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.messages.Message;
|
||||
import org.springframework.ai.chat.messages.MessageType;
|
||||
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
|
||||
import org.springframework.ai.chat.metadata.DefaultUsage;
|
||||
import org.springframework.ai.chat.metadata.RateLimit;
|
||||
import org.springframework.ai.chat.model.ChatModel;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.model.Generation;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.ai.chat.prompt.PromptTemplate;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.mockito.BDDMockito.given;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @author Timo Salm
|
||||
* @author Alexandros Pappas
|
||||
*/
|
||||
@ExtendWith(MockitoExtension.class)
|
||||
public class QuestionAnswerAdvisorTests {
|
||||
|
||||
@Mock
|
||||
ChatModel chatModel;
|
||||
|
||||
@Captor
|
||||
ArgumentCaptor<Prompt> promptCaptor;
|
||||
|
||||
@Captor
|
||||
ArgumentCaptor<SearchRequest> vectorSearchCaptor;
|
||||
|
||||
@Mock
|
||||
VectorStore vectorStore;
|
||||
|
||||
@Test
|
||||
public void qaAdvisorWithDynamicFilterExpressions() {
|
||||
|
||||
// @formatter:off
|
||||
given(this.chatModel.call(this.promptCaptor.capture()))
|
||||
.willReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("Your answer is ZXY"))),
|
||||
ChatResponseMetadata.builder().id("678").model("model1").keyValue("key6", "value6").metadata(Map.of("key1", "value1")).promptMetadata(null).rateLimit(new RateLimit() {
|
||||
|
||||
@Override
|
||||
public Long getRequestsLimit() {
|
||||
return 5L;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Long getRequestsRemaining() {
|
||||
return 6L;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Duration getRequestsReset() {
|
||||
return Duration.ofSeconds(7);
|
||||
}
|
||||
|
||||
@Override
|
||||
public Long getTokensLimit() {
|
||||
return 8L;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Long getTokensRemaining() {
|
||||
return 8L;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Duration getTokensReset() {
|
||||
return Duration.ofSeconds(9);
|
||||
}
|
||||
}).usage(new DefaultUsage(6L, 7L))
|
||||
.build()));
|
||||
// @formatter:on
|
||||
|
||||
given(this.vectorStore.similaritySearch(this.vectorSearchCaptor.capture()))
|
||||
.willReturn(List.of(new Document("doc1"), new Document("doc2")));
|
||||
|
||||
var qaAdvisor = new QuestionAnswerAdvisor(this.vectorStore,
|
||||
SearchRequest.defaults().withSimilarityThreshold(0.99d).withTopK(6));
|
||||
|
||||
var chatClient = ChatClient.builder(this.chatModel)
|
||||
.defaultSystem("Default system text.")
|
||||
.defaultAdvisors(qaAdvisor)
|
||||
.build();
|
||||
|
||||
// @formatter:off
|
||||
var response = chatClient.prompt()
|
||||
.user("Please answer my question XYZ")
|
||||
.advisors(a -> a.param(QuestionAnswerAdvisor.FILTER_EXPRESSION, "type == 'Spring'"))
|
||||
.call()
|
||||
.chatResponse();
|
||||
//formatter:on
|
||||
|
||||
// Ensure the metadata is correctly copied over
|
||||
assertThat(response.getMetadata().getModel()).isEqualTo("model1");
|
||||
assertThat(response.getMetadata().getId()).isEqualTo("678");
|
||||
assertThat(response.getMetadata().getRateLimit().getRequestsLimit()).isEqualTo(5L);
|
||||
assertThat(response.getMetadata().getRateLimit().getRequestsRemaining()).isEqualTo(6L);
|
||||
assertThat(response.getMetadata().getRateLimit().getRequestsReset()).isEqualTo(Duration.ofSeconds(7));
|
||||
assertThat(response.getMetadata().getRateLimit().getTokensLimit()).isEqualTo(8L);
|
||||
assertThat(response.getMetadata().getRateLimit().getTokensRemaining()).isEqualTo(8L);
|
||||
assertThat(response.getMetadata().getRateLimit().getTokensReset()).isEqualTo(Duration.ofSeconds(9));
|
||||
assertThat(response.getMetadata().getUsage().getPromptTokens()).isEqualTo(6L);
|
||||
assertThat(response.getMetadata().getUsage().getGenerationTokens()).isEqualTo(7L);
|
||||
assertThat(response.getMetadata().getUsage().getTotalTokens()).isEqualTo(6L + 7L);
|
||||
assertThat(response.getMetadata().get("key6").toString()).isEqualTo("value6");
|
||||
assertThat(response.getMetadata().get("key1").toString()).isEqualTo("value1");
|
||||
|
||||
String content = response.getResult().getOutput().getText();
|
||||
|
||||
assertThat(content).isEqualTo("Your answer is ZXY");
|
||||
|
||||
Message systemMessage = this.promptCaptor.getValue().getInstructions().get(0);
|
||||
|
||||
System.out.println(systemMessage.getText());
|
||||
|
||||
assertThat(systemMessage.getText()).isEqualToIgnoringWhitespace("""
|
||||
Default system text.
|
||||
""");
|
||||
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
|
||||
|
||||
Message userMessage = this.promptCaptor.getValue().getInstructions().get(1);
|
||||
|
||||
assertThat(userMessage.getText()).isEqualToIgnoringWhitespace("""
|
||||
Please answer my question XYZ
|
||||
Context information is below, surrounded by ---------------------
|
||||
|
||||
---------------------
|
||||
doc1
|
||||
doc2
|
||||
---------------------
|
||||
|
||||
Given the context and provided history information and not prior knowledge,
|
||||
reply to the user comment. If the answer is not in the context, inform
|
||||
the user that you can't answer the question.
|
||||
""");
|
||||
|
||||
assertThat(this.vectorSearchCaptor.getValue().getFilterExpression()).isEqualTo(new FilterExpressionBuilder().eq("type", "Spring").build());
|
||||
assertThat(this.vectorSearchCaptor.getValue().getSimilarityThreshold()).isEqualTo(0.99d);
|
||||
assertThat(this.vectorSearchCaptor.getValue().getTopK()).isEqualTo(6);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void qaAdvisorTakesUserTextParametersIntoAccountForSimilaritySearch() {
|
||||
given(this.chatModel.call(this.promptCaptor.capture()))
|
||||
.willReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("Your answer is ZXY"))),
|
||||
ChatResponseMetadata.builder().build()));
|
||||
|
||||
given(this.vectorStore.similaritySearch(this.vectorSearchCaptor.capture()))
|
||||
.willReturn(List.of(new Document("doc1"), new Document("doc2")));
|
||||
|
||||
var chatClient = ChatClient.builder(this.chatModel).build();
|
||||
var qaAdvisor = new QuestionAnswerAdvisor(this.vectorStore, SearchRequest.defaults());
|
||||
|
||||
var userTextTemplate = "Please answer my question {question}";
|
||||
// @formatter:off
|
||||
chatClient.prompt()
|
||||
.user(u -> u.text(userTextTemplate).param("question", "XYZ"))
|
||||
.advisors(qaAdvisor)
|
||||
.call()
|
||||
.chatResponse();
|
||||
//formatter:on
|
||||
|
||||
var expectedQuery = "Please answer my question XYZ";
|
||||
var userPrompt = this.promptCaptor.getValue().getInstructions().get(0).getText();
|
||||
assertThat(userPrompt).doesNotContain(userTextTemplate);
|
||||
assertThat(userPrompt).contains(expectedQuery);
|
||||
assertThat(this.vectorSearchCaptor.getValue().getQuery()).isEqualTo(expectedQuery);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void qaAdvisorTakesUserParameterizedUserMessagesIntoAccountForSimilaritySearch() {
|
||||
given(this.chatModel.call(this.promptCaptor.capture()))
|
||||
.willReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("Your answer is ZXY"))),
|
||||
ChatResponseMetadata.builder().build()));
|
||||
|
||||
given(this.vectorStore.similaritySearch(this.vectorSearchCaptor.capture()))
|
||||
.willReturn(List.of(new Document("doc1"), new Document("doc2")));
|
||||
|
||||
var chatClient = ChatClient.builder(this.chatModel).build();
|
||||
var qaAdvisor = new QuestionAnswerAdvisor(this.vectorStore, SearchRequest.defaults());
|
||||
|
||||
var userTextTemplate = "Please answer my question {question}";
|
||||
var userPromptTemplate = new PromptTemplate(userTextTemplate, Map.of("question", "XYZ"));
|
||||
var userMessage = userPromptTemplate.createMessage();
|
||||
// @formatter:off
|
||||
chatClient.prompt(new Prompt(userMessage))
|
||||
.advisors(qaAdvisor)
|
||||
.call()
|
||||
.chatResponse();
|
||||
//formatter:on
|
||||
|
||||
var expectedQuery = "Please answer my question XYZ";
|
||||
var userPrompt = this.promptCaptor.getValue().getInstructions().get(0).getText();
|
||||
assertThat(userPrompt).doesNotContain(userTextTemplate);
|
||||
assertThat(userPrompt).contains(expectedQuery);
|
||||
assertThat(this.vectorSearchCaptor.getValue().getQuery()).isEqualTo(expectedQuery);
|
||||
}
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user