From c4784a070f13e0f8714e5ac71f358980bcdb3e82 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Sat, 25 May 2024 13:17:29 +0200 Subject: [PATCH] Add ChatClient plugable advisors support - Add RequestResponseAdvisor interface with adviseReqeust and adviseResponse methods. The adviseRequest method takes and returns AdvisedRequest. The adviseResponse method takes and returns ChatRequest. - Add ChatClient#ChatClientRequets advisor(...) methods to register advisros. ChatClient call the registered advisors in order before sealing the ChatClientRequest into a Prompt and call the model and after the model response. - Implement PromptChatMemoryAdvisor that uses the ChatMemory and the systemem prompt. - Implement MessageChatMemoryAdvisor that usees the ChatMemory and the prompt messages. - Implement VectorStoreChatMemoryAdvisor that uses VecrStore for long term message history. - Add tests. - Add shared context to the RequestResponseAdvisor's flow, Context is shared between the request and the response - Add advisor parameters that are passed through the context --- .../ai/chat/client/AdvisedRequest.java | 156 ++++++++++ .../ai/chat/client/ChatClient.java | 287 ++++++++++++++---- .../chat/client/RequestResponseAdvisor.java | 73 +++++ .../advisor/AbstractChatMemoryAdvisor.java | 77 +++++ .../LastMaxTokenSizeContentPurger.java | 74 +++++ .../advisor/MessageChatMemoryAdvisor.java | 95 ++++++ .../advisor/PromptChatMemoryAdvisor.java | 124 ++++++++ .../client/advisor/QuestionAnswerAdvisor.java | 98 ++++++ .../advisor/VectorStoreChatMemoryAdvisor.java | 145 +++++++++ .../memory/ChatMemoryChatServiceListener.java | 5 + .../ai/chat/memory/InMemoryChatMemory.java | 9 + .../LastMaxTokenSizeContentTransformer.java | 4 + .../memory/MessageChatMemoryAugmentor.java | 3 + .../SystemPromptChatMemoryAugmentor.java | 3 + ...torStoreChatMemoryChatServiceListener.java | 3 + .../VectorStoreChatMemoryRetriever.java | 3 + .../ai/chat/model/ChatModel.java | 8 +- .../{service => model}/MessageAggregator.java | 5 +- ...treamingPromptTransformingChatService.java | 1 + .../chat/client/ChatClientAdvisorTests.java | 272 +++++++++++++++++ .../ai/chat/client/ChatClientTest.java | 13 +- 21 files changed, 1378 insertions(+), 80 deletions(-) create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/client/AdvisedRequest.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/client/RequestResponseAdvisor.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/AbstractChatMemoryAdvisor.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/LastMaxTokenSizeContentPurger.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/MessageChatMemoryAdvisor.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/PromptChatMemoryAdvisor.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/QuestionAnswerAdvisor.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/VectorStoreChatMemoryAdvisor.java rename spring-ai-core/src/main/java/org/springframework/ai/chat/{service => model}/MessageAggregator.java (94%) create mode 100644 spring-ai-core/src/test/java/org/springframework/ai/chat/client/ChatClientAdvisorTests.java diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/AdvisedRequest.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/AdvisedRequest.java new file mode 100644 index 000000000..16337a33c --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/AdvisedRequest.java @@ -0,0 +1,156 @@ +/* + * Copyright 2024-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; + +import java.util.List; +import java.util.Map; + +import org.springframework.ai.chat.messages.Media; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.model.ChatModel; +import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.model.function.FunctionCallback; + +/** + * The data of the chat client request that can be modifed before the execution of the + * ChatClient's call method + * + * @author Christian Tzolov + * @since 1.0.0 M1 + * + */ +public record AdvisedRequest(ChatModel chatModel, String userText, String systemText, ChatOptions chatOptions, + List media, List functionNames, List functionCallbacks, List messages, + Map userParams, Map systemParams, List advisors, + Map advisorParams) { + + public static Builder from(AdvisedRequest from) { + Builder builder = new Builder(); + builder.chatModel = from.chatModel; + builder.userText = from.userText; + builder.systemText = from.systemText; + builder.chatOptions = from.chatOptions; + builder.media = from.media; + builder.functionNames = from.functionNames; + builder.functionCallbacks = from.functionCallbacks; + builder.messages = from.messages; + builder.userParams = from.userParams; + builder.systemParams = from.systemParams; + builder.advisors = from.advisors; + builder.advisorParams = from.advisorParams; + return builder; + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + + private ChatModel chatModel; + + private String userText = ""; + + private String systemText = ""; + + private ChatOptions chatOptions = null; + + private List media = List.of(); + + private List functionNames = List.of(); + + private List functionCallbacks = List.of(); + + private List messages = List.of(); + + private Map userParams = Map.of(); + + private Map systemParams = Map.of(); + + private List advisors = List.of(); + + private Map advisorParams = Map.of(); + + public Builder withChatModel(ChatModel chatModel) { + this.chatModel = chatModel; + return this; + } + + public Builder withUserText(String userText) { + this.userText = userText; + return this; + } + + public Builder withSystemText(String systemText) { + this.systemText = systemText; + return this; + } + + public Builder withChatOptions(ChatOptions chatOptions) { + this.chatOptions = chatOptions; + return this; + } + + public Builder withMedia(List media) { + this.media = media; + return this; + } + + public Builder withFunctionNames(List functionNames) { + this.functionNames = functionNames; + return this; + } + + public Builder withFunctionCallbacks(List functionCallbacks) { + this.functionCallbacks = functionCallbacks; + return this; + } + + public Builder withMessages(List messages) { + this.messages = messages; + return this; + } + + public Builder withUserParams(Map userParams) { + this.userParams = userParams; + return this; + } + + public Builder withSystemParams(Map systemParams) { + this.systemParams = systemParams; + return this; + } + + public Builder withAdvisors(List advisors) { + this.advisors = advisors; + return this; + } + + public Builder withAdvisorParams(Map advisorParams) { + this.advisorParams = advisorParams; + return this; + } + + public AdvisedRequest build() { + return new AdvisedRequest(chatModel, this.userText, this.systemText, this.chatOptions, this.media, + this.functionNames, this.functionCallbacks, this.messages, this.userParams, this.systemParams, + this.advisors, this.advisorParams); + } + + } + +} \ No newline at end of file diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/ChatClient.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/ChatClient.java index aacdb25d3..17709827e 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/ChatClient.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/ChatClient.java @@ -24,6 +24,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import reactor.core.publisher.Flux; @@ -208,6 +209,34 @@ public interface ChatClient { } + class AdvisorSpec { + + private List advisors = new ArrayList<>(); + + private final Map params = new HashMap<>(); + + public AdvisorSpec param(String k, Object v) { + this.params.put(k, v); + return this; + } + + public AdvisorSpec params(Map p) { + this.params.putAll(p); + return this; + } + + public AdvisorSpec advisors(RequestResponseAdvisor... advisors) { + this.advisors.addAll(List.of(advisors)); + return this; + } + + public AdvisorSpec advisors(List advisors) { + this.advisors.addAll(advisors); + return this; + } + + } + class ChatClientRequest { private final ChatModel chatModel; @@ -230,6 +259,37 @@ public interface ChatClient { private final Map systemParams = new HashMap<>(); + private List advisors = new ArrayList<>(); + + private final Map advisorParams = new HashMap<>(); + + /* copy constructor */ + ChatClientRequest(ChatClientRequest ccr) { + this(ccr.chatModel, ccr.userText, ccr.userParams, ccr.systemText, ccr.systemParams, ccr.functionCallbacks, + ccr.messages, ccr.functionNames, ccr.media, ccr.chatOptions, ccr.advisors, ccr.advisorParams); + } + + public ChatClientRequest(ChatModel chatModel, String userText, Map userParams, + String systemText, Map systemParams, List functionCallbacks, + List messages, List functionNames, List media, ChatOptions chatOptions, + List advisors, Map advisorParams) { + + this.chatModel = chatModel; + this.chatOptions = chatOptions != null ? chatOptions : chatModel.getDefaultOptions(); + + this.userText = userText; + this.userParams.putAll(userParams); + this.systemText = systemText; + this.systemParams.putAll(systemParams); + + this.functionNames.addAll(functionNames); + this.functionCallbacks.addAll(functionCallbacks); + this.messages.addAll(messages); + this.media.addAll(media); + this.advisors.addAll(advisors); + this.advisorParams.putAll(advisorParams); + } + /** * Return a {@code ChatClient.Builder} to create a new {@code ChatClient} whose * settings are replicated from this {@code ChatClientRequest}. @@ -250,46 +310,52 @@ public interface ChatClient { return builder; } - /* copy constructor */ - ChatClientRequest(ChatClientRequest ccr) { - this(ccr.chatModel, ccr.userText, ccr.userParams, ccr.systemText, ccr.systemParams, ccr.functionCallbacks, - ccr.functionNames, ccr.media, ccr.chatOptions); + public ChatClientRequest advisors(Consumer consumer) { + Assert.notNull(consumer, "the consumer must be non-null"); + var as = new AdvisorSpec(); + consumer.accept(as); + this.advisorParams.putAll(as.params); + this.advisors.addAll(as.advisors); + return this; } - public ChatClientRequest(ChatModel chatModel, String userText, Map userParams, - String systemText, Map systemParams, List functionCallbacks, - List functionNames, List media, ChatOptions chatOptions) { + public ChatClientRequest advisors(RequestResponseAdvisor... advisors) { + Assert.notNull(advisors, "the advisors must be non-null"); + this.advisors.addAll(List.of(advisors)); + return this; + } - this.chatModel = chatModel; - this.chatOptions = chatOptions != null ? chatOptions : chatModel.getDefaultOptions(); - - this.userText = userText; - this.userParams.putAll(userParams); - this.systemText = systemText; - this.systemParams.putAll(systemParams); - - this.functionNames.addAll(functionNames); - this.functionCallbacks.addAll(functionCallbacks); - this.media.addAll(media); + public ChatClientRequest advisors(List advisors) { + Assert.notNull(advisors, "the advisors must be non-null"); + this.advisors.addAll(advisors); + return this; } public ChatClientRequest messages(Message... messages) { + Assert.notNull(messages, "the messages must be non-null"); this.messages.addAll(List.of(messages)); return this; } public ChatClientRequest messages(List messages) { + Assert.notNull(messages, "the messages must be non-null"); this.messages.addAll(messages); return this; } public ChatClientRequest options(T options) { + Assert.notNull(options, "the options must be non-null"); this.chatOptions = options; return this; } public ChatClientRequest function(String name, String description, java.util.function.Function function) { + + Assert.hasText(name, "the name must be non-null and non-empty"); + Assert.hasText(description, "the description must be non-null and non-empty"); + Assert.notNull(function, "the function must be non-null"); + var fcw = FunctionCallbackWrapper.builder(function) .withDescription(description) .withName(name) @@ -300,18 +366,24 @@ public interface ChatClient { } public ChatClientRequest functions(String... functionBeanNames) { + Assert.notNull(functionBeanNames, "the functionBeanNames must be non-null"); this.functionNames.addAll(List.of(functionBeanNames)); return this; } public ChatClientRequest system(String text) { + Assert.notNull(text, "the text must be non-null"); this.systemText = text; return this; } - public ChatClientRequest system(Resource text, Charset charset) { + public ChatClientRequest system(Resource textResource, Charset charset) { + + Assert.notNull(textResource, "the text resource must be non-null"); + Assert.notNull(charset, "the charset must be non-null"); + try { - this.systemText = text.getContentAsString(charset); + this.systemText = textResource.getContentAsString(charset); } catch (IOException e) { throw new RuntimeException(e); @@ -320,10 +392,14 @@ public interface ChatClient { } public ChatClientRequest system(Resource text) { + Assert.notNull(text, "the text resource must be non-null"); return this.system(text, Charset.defaultCharset()); } public ChatClientRequest system(Consumer consumer) { + + Assert.notNull(consumer, "the consumer must be non-null"); + var ss = new SystemSpec(); consumer.accept(ss); this.systemText = StringUtils.hasText(ss.text()) ? ss.text() : this.systemText; @@ -333,11 +409,16 @@ public interface ChatClient { } public ChatClientRequest user(String text) { + Assert.notNull(text, "the text must be non-null"); this.userText = text; return this; } public ChatClientRequest user(Resource text, Charset charset) { + + Assert.notNull(text, "the text resource must be non-null"); + Assert.notNull(charset, "the charset must be non-null"); + try { this.userText = text.getContentAsString(charset); } @@ -348,10 +429,13 @@ public interface ChatClient { } public ChatClientRequest user(Resource text) { + Assert.notNull(text, "the text resource must be non-null"); return this.user(text, Charset.defaultCharset()); } public ChatClientRequest user(Consumer consumer) { + Assert.notNull(consumer, "the consumer must be non-null"); + var us = new UserSpec(); consumer.accept(us); this.userText = StringUtils.hasText(us.text()) ? us.text() : this.userText; @@ -423,6 +507,33 @@ public interface ChatClient { } + private static ChatClientRequest adviseOnRequest(ChatClientRequest inputRequest, Map context) { + + ChatClientRequest advisedRequest = inputRequest; + + if (!CollectionUtils.isEmpty(inputRequest.advisors)) { + AdvisedRequest adviseRequest = new AdvisedRequest(inputRequest.chatModel, inputRequest.userText, + inputRequest.systemText, inputRequest.chatOptions, inputRequest.media, + inputRequest.functionNames, inputRequest.functionCallbacks, inputRequest.messages, + inputRequest.userParams, inputRequest.systemParams, inputRequest.advisors, + inputRequest.advisorParams); + + // apply the advisors onRequest + var currentAdvisors = new ArrayList<>(inputRequest.advisors); + for (RequestResponseAdvisor advisor : currentAdvisors) { + adviseRequest = advisor.adviseRequest(adviseRequest, context); + } + + advisedRequest = new ChatClientRequest(adviseRequest.chatModel(), adviseRequest.userText(), + adviseRequest.userParams(), adviseRequest.systemText(), adviseRequest.systemParams(), + adviseRequest.functionCallbacks(), adviseRequest.messages(), adviseRequest.functionNames(), + adviseRequest.media(), adviseRequest.chatOptions(), adviseRequest.advisors(), + adviseRequest.advisorParams()); + } + + return advisedRequest; + } + public static class CallResponseSpec { private final ChatClientRequest request; @@ -443,7 +554,7 @@ public interface ChatClient { } private T doSingleWithBeanOutputConverter(StructuredOutputConverter boc) { - var chatResponse = doGetChatResponse(boc.getFormat()); + var chatResponse = doGetChatResponse(this.request, boc.getFormat()); var stringResponse = chatResponse.getResult().getOutput().getContent(); return boc.convert(stringResponse); } @@ -455,48 +566,65 @@ public interface ChatClient { } private ChatResponse doGetChatResponse() { - return this.doGetChatResponse(""); + return this.doGetChatResponse(this.request, ""); } - private ChatResponse doGetChatResponse(String formatParam) { + private ChatResponse doGetChatResponse(ChatClientRequest inputRequest, String formatParam) { + + Map context = new ConcurrentHashMap<>(); + context.putAll(inputRequest.advisorParams); + ChatClientRequest advisedRequest = adviseOnRequest(inputRequest, context); var processedUserText = StringUtils.hasText(formatParam) - ? this.request.userText + System.lineSeparator() + "{format}" : this.request.userText; + ? advisedRequest.userText + System.lineSeparator() + "{spring.ai.soc.format}" + : advisedRequest.userText; - Map userParams = new HashMap<>(this.request.userParams); + Map userParams = new HashMap<>(advisedRequest.userParams); if (StringUtils.hasText(formatParam)) { - userParams.put("format", formatParam); + userParams.put("spring.ai.soc.format", formatParam); } - var messages = new ArrayList(this.request.messages); + var messages = new ArrayList(advisedRequest.messages); var textsAreValid = (StringUtils.hasText(processedUserText) - || StringUtils.hasText(this.request.systemText)); + || StringUtils.hasText(advisedRequest.systemText)); if (textsAreValid) { - if (StringUtils.hasText(this.request.systemText) || !this.request.systemParams.isEmpty()) { + if (StringUtils.hasText(advisedRequest.systemText) || !advisedRequest.systemParams.isEmpty()) { var systemMessage = new SystemMessage( - new PromptTemplate(this.request.systemText, this.request.systemParams).render()); + new PromptTemplate(advisedRequest.systemText, advisedRequest.systemParams).render()); messages.add(systemMessage); } UserMessage userMessage = null; if (!CollectionUtils.isEmpty(userParams)) { userMessage = new UserMessage(new PromptTemplate(processedUserText, userParams).render(), - this.request.media); + advisedRequest.media); } else { - userMessage = new UserMessage(processedUserText, this.request.media); + userMessage = new UserMessage(processedUserText, advisedRequest.media); } messages.add(userMessage); } - if (this.request.chatOptions instanceof FunctionCallingOptions functionCallingOptions) { - if (!this.request.functionNames.isEmpty()) { - functionCallingOptions.setFunctions(new HashSet<>(this.request.functionNames)); + + if (advisedRequest.chatOptions instanceof FunctionCallingOptions functionCallingOptions) { + if (!advisedRequest.functionNames.isEmpty()) { + functionCallingOptions.setFunctions(new HashSet<>(advisedRequest.functionNames)); } - if (!this.request.functionCallbacks.isEmpty()) { - functionCallingOptions.setFunctionCallbacks(this.request.functionCallbacks); + if (!advisedRequest.functionCallbacks.isEmpty()) { + functionCallingOptions.setFunctionCallbacks(advisedRequest.functionCallbacks); } } - var prompt = new Prompt(messages, this.request.chatOptions); - return this.chatModel.call(prompt); + var prompt = new Prompt(messages, advisedRequest.chatOptions); + var chatResponse = this.chatModel.call(prompt); + + ChatResponse advisedResponse = chatResponse; + // apply the advisors on response + if (!CollectionUtils.isEmpty(inputRequest.advisors)) { + var currentAdvisors = new ArrayList<>(inputRequest.advisors); + for (RequestResponseAdvisor advisor : currentAdvisors) { + advisedResponse = advisor.adviseResponse(advisedResponse, context); + } + } + + return advisedResponse; } public ChatResponse chatResponse() { @@ -520,55 +648,67 @@ public interface ChatClient { this.request = request; } - private Flux doGetFluxChatResponse(String processedUserText) { - Map userParams = new HashMap<>(this.request.userParams); + private Flux doGetFluxChatResponse(ChatClientRequest inputRequest) { - var messages = new ArrayList(); + Map context = new ConcurrentHashMap<>(); + context.putAll(inputRequest.advisorParams); + ChatClientRequest advisedRequest = adviseOnRequest(inputRequest, context); + + String processedUserText = advisedRequest.userText; + Map userParams = new HashMap<>(advisedRequest.userParams); + + var messages = new ArrayList(advisedRequest.messages); var textsAreValid = (StringUtils.hasText(processedUserText) - || StringUtils.hasText(this.request.systemText)); - var messagesAreValid = !this.request.messages.isEmpty(); - Assert.state(!(messagesAreValid && textsAreValid), "you must specify either " + Message.class.getName() - + " instances or user/system texts, but not both"); + || StringUtils.hasText(advisedRequest.systemText)); if (textsAreValid) { UserMessage userMessage = null; if (!CollectionUtils.isEmpty(userParams)) { userMessage = new UserMessage(new PromptTemplate(processedUserText, userParams).render(), - this.request.media); + advisedRequest.media); } else { - userMessage = new UserMessage(processedUserText, this.request.media); + userMessage = new UserMessage(processedUserText, advisedRequest.media); } - if (StringUtils.hasText(this.request.systemText) || !this.request.systemParams.isEmpty()) { + if (StringUtils.hasText(advisedRequest.systemText) || !advisedRequest.systemParams.isEmpty()) { var systemMessage = new SystemMessage( - new PromptTemplate(this.request.systemText, this.request.systemParams).render()); + new PromptTemplate(advisedRequest.systemText, advisedRequest.systemParams).render()); messages.add(systemMessage); } messages.add(userMessage); } - else { - messages.addAll(this.request.messages); - } - if (this.request.chatOptions instanceof FunctionCallingOptions functionCallingOptions) { - // if (this.request.chatOptions instanceof - // FunctionCallingOptionsBuilder.PortableFunctionCallingOptions - // functionCallingOptions) { - if (!this.request.functionNames.isEmpty()) { - functionCallingOptions.setFunctions(new HashSet<>(this.request.functionNames)); + + if (advisedRequest.chatOptions instanceof + + FunctionCallingOptions functionCallingOptions) { + if (!advisedRequest.functionNames.isEmpty()) { + functionCallingOptions.setFunctions(new HashSet<>(advisedRequest.functionNames)); } - if (!this.request.functionCallbacks.isEmpty()) { - functionCallingOptions.setFunctionCallbacks(this.request.functionCallbacks); + if (!advisedRequest.functionCallbacks.isEmpty()) { + functionCallingOptions.setFunctionCallbacks(advisedRequest.functionCallbacks); } } - var prompt = new Prompt(messages, this.request.chatOptions); - return this.chatModel.stream(prompt); + var prompt = new Prompt(messages, advisedRequest.chatOptions); + + var fluxChatResponse = this.chatModel.stream(prompt); + + Flux advisedResponse = fluxChatResponse; + // apply the advisors on response + if (!CollectionUtils.isEmpty(inputRequest.advisors)) { + var currentAdvisors = new ArrayList<>(inputRequest.advisors); + for (RequestResponseAdvisor advisor : currentAdvisors) { + advisedResponse = advisor.adviseResponse(advisedResponse, context); + } + } + + return advisedResponse; } public Flux chatResponse() { - return doGetFluxChatResponse(this.request.userText); + return doGetFluxChatResponse(this.request); } public Flux content() { - return doGetFluxChatResponse(this.request.userText).map(r -> { + return doGetFluxChatResponse(this.request).map(r -> { if (r.getResult() == null || r.getResult().getOutput() == null || r.getResult().getOutput().getContent() == null) { return ""; @@ -599,7 +739,22 @@ public interface ChatClient { Assert.notNull(chatModel, "the " + ChatModel.class.getName() + " must be non-null"); this.chatModel = chatModel; this.defaultRequest = new ChatClientRequest(chatModel, "", Map.of(), "", Map.of(), List.of(), List.of(), - List.of(), null); + List.of(), List.of(), null, List.of(), Map.of()); + } + + public Builder defaultAdvisors(RequestResponseAdvisor advisor) { + this.defaultRequest.advisors(advisor); + return this; + } + + public Builder defaultAdvisors(Consumer advisorSpecConsumer) { + this.defaultRequest.advisors(advisorSpecConsumer); + return this; + } + + public Builder defaultAdvisors(List advisors) { + this.defaultRequest.advisors(advisors); + return this; } public ChatClient build() { diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/RequestResponseAdvisor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/RequestResponseAdvisor.java new file mode 100644 index 000000000..922898f65 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/RequestResponseAdvisor.java @@ -0,0 +1,73 @@ +/* + * Copyright 2024-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; + +import java.util.Map; + +import reactor.core.publisher.Flux; + +import org.springframework.ai.chat.client.ChatClient.ChatClientRequest; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.ChatModel; +import org.springframework.ai.chat.prompt.Prompt; + +/** + * Advisor called before and after the {@link ChatModel#call(Prompt)} and + * {@link ChatModel#stream(Prompt)} methods calls. The {@link ChatClient} maintains a + * chain of advisors with chared execution context. + * + * @author Christian Tzolov + * @since 1.0.0 M1 + */ +public interface RequestResponseAdvisor { + + /** + * @param request the {@link AdvisedRequest} data to be advised. Represents the row + * {@link ChatClientRequest} data before sealed into a {@link Prompt}. + * @param context the shared data between the advisors in the chain. It is shared + * between all request and response advising points of all advisors in the chain. + * @return the advised {@link AdvisedRequest}. + */ + default AdvisedRequest adviseRequest(AdvisedRequest request, Map context) { + return request; + } + + /** + * @param response the {@link ChatResponse} data to be advised. Represents the row + * {@link ChatResponse} data after the {@link ChatModel#call(Prompt)} method is + * called. + * @param context the shared data between the advisors in the chain. It is shared + * between all request and response advising points of all advisors in the chain. + * @return the advised {@link ChatResponse}. + */ + default ChatResponse adviseResponse(ChatResponse response, Map context) { + return response; + } + + /** + * @param fluxResponse the streaming {@link ChatResponse} data to be advised. + * Represents the row {@link ChatResponse} stream data after the + * {@link ChatModel#stream(Prompt)} method is called. + * @param context the shared data between the advisors in the chain. It is shared + * between all request and response advising points of all advisors in the chain. + * @return the advised {@link ChatResponse} flux. + */ + default Flux adviseResponse(Flux fluxResponse, Map context) { + return fluxResponse; + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/AbstractChatMemoryAdvisor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/AbstractChatMemoryAdvisor.java new file mode 100644 index 000000000..edef56e7d --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/AbstractChatMemoryAdvisor.java @@ -0,0 +1,77 @@ +/* + * Copyright 2024-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; + +import java.util.Map; + +import org.springframework.ai.chat.client.RequestResponseAdvisor; +import org.springframework.util.Assert; + +/** + * Abstract class that serves as a base for chat memory advisors. + * + * @author Christian Tzolov + * @since 1.0.0 M1 + */ +public abstract class AbstractChatMemoryAdvisor implements RequestResponseAdvisor { + + public static final String CHAT_MEMORY_CONVERSATION_ID_KEY = "chat.memory.conversation.id"; + + public static final String CHAT_MEMORY_RETRIEVE_SIZE_KEY = "chat.memory.response.size"; + + public static final String DEFAULT_CHAT_MEMORY_CONVERSATION_ID = "default"; + + public static final int DEFAULT_CHAT_MEMORY_RESPONSE_SIZE = 100; + + protected final T chatMemoryStore; + + protected final String defaultConversationId; + + protected final int defaultChatMemoryRetrieveSize; + + public AbstractChatMemoryAdvisor(T chatMemory) { + this(chatMemory, DEFAULT_CHAT_MEMORY_CONVERSATION_ID, DEFAULT_CHAT_MEMORY_RESPONSE_SIZE); + } + + public AbstractChatMemoryAdvisor(T chatMemory, String defaultConversationId, int defaultChatMemoryRetrieveSize) { + + Assert.notNull(chatMemory, "The chatMemory must not be null!"); + Assert.hasText(defaultConversationId, "The conversationId must not be empty!"); + Assert.isTrue(defaultChatMemoryRetrieveSize > 0, "The defaultChatMemoryRetrieveSize must be greater than 0!"); + + this.chatMemoryStore = chatMemory; + this.defaultConversationId = defaultConversationId; + this.defaultChatMemoryRetrieveSize = defaultChatMemoryRetrieveSize; + } + + protected T getChatMemoryStore() { + return this.chatMemoryStore; + } + + protected String doGetConversationId(Map context) { + + return context.containsKey(CHAT_MEMORY_CONVERSATION_ID_KEY) + ? context.get(CHAT_MEMORY_CONVERSATION_ID_KEY).toString() : this.defaultConversationId; + } + + protected int doGetChatMemoryRetrieveSize(Map context) { + return context.containsKey(CHAT_MEMORY_RETRIEVE_SIZE_KEY) + ? Integer.parseInt(context.get(CHAT_MEMORY_RETRIEVE_SIZE_KEY).toString()) + : this.defaultChatMemoryRetrieveSize; + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/LastMaxTokenSizeContentPurger.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/LastMaxTokenSizeContentPurger.java new file mode 100644 index 000000000..ac9210234 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/LastMaxTokenSizeContentPurger.java @@ -0,0 +1,74 @@ +/* + * Copyright 2024-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; + +import java.util.ArrayList; +import java.util.List; + +import org.springframework.ai.model.Content; +import org.springframework.ai.tokenizer.TokenCountEstimator; + +/** + * Returns a new list of content (e.g list of messages of list of documents) that is a + * subset of the input list of contents and complies with the max token size constraint. + * + * The token estimator is used to estimate the token count of the datum. + * + * @author Christian Tzolov + * @since 1.0.0 M1 + */ +public class LastMaxTokenSizeContentPurger { + + protected final TokenCountEstimator tokenCountEstimator; + + protected final int maxTokenSize; + + public LastMaxTokenSizeContentPurger(TokenCountEstimator tokenCountEstimator, int maxTokenSize) { + this.tokenCountEstimator = tokenCountEstimator; + this.maxTokenSize = maxTokenSize; + } + + public List purgeExcess(List datum, int totalSize) { + + int index = 0; + List newList = new ArrayList<>(); + + while (index < datum.size() && totalSize > this.maxTokenSize) { + Content oldDatum = datum.get(index++); + int oldMessageTokenSize = this.doEstimateTokenCount(oldDatum); + totalSize = totalSize - oldMessageTokenSize; + } + + if (index >= datum.size()) { + return List.of(); + } + + // add the rest of the messages. + newList.addAll(datum.subList(index, datum.size())); + + return newList; + } + + protected int doEstimateTokenCount(Content datum) { + return this.tokenCountEstimator.estimate(datum); + } + + protected int doEstimateTokenCount(List datum) { + return datum.stream().mapToInt(this::doEstimateTokenCount).sum(); + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/MessageChatMemoryAdvisor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/MessageChatMemoryAdvisor.java new file mode 100644 index 000000000..bf22231ce --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/MessageChatMemoryAdvisor.java @@ -0,0 +1,95 @@ +/* + * Copyright 2024-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; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +import reactor.core.publisher.Flux; + +import org.springframework.ai.chat.client.AdvisedRequest; +import org.springframework.ai.chat.memory.ChatMemory; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.MessageAggregator; + +/** + * Memory is retrieved added as a collection of messages to the prompt + * + * @author Christian Tzolov + * @since 1.0.0 M1 + */ +public class MessageChatMemoryAdvisor extends AbstractChatMemoryAdvisor { + + public MessageChatMemoryAdvisor(ChatMemory chatMemory) { + super(chatMemory); + } + + public MessageChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int chatHistoryWindowSize) { + super(chatMemory, defaultConversationId, chatHistoryWindowSize); + } + + @Override + public AdvisedRequest adviseRequest(AdvisedRequest request, Map context) { + + String conversationId = this.doGetConversationId(context); + + int chatMemoryRetrieveSize = this.doGetChatMemoryRetrieveSize(context); + + // 1. Retrieve the chat memory for the current conversation. + List memoryMessages = this.getChatMemoryStore().get(conversationId, chatMemoryRetrieveSize); + + // 2. Advise the request messages list. + List advisedMessages = new ArrayList<>(request.messages()); + advisedMessages.addAll(memoryMessages); + + // 3. Create a new request with the advised messages. + AdvisedRequest advisedRequest = AdvisedRequest.from(request).withMessages(advisedMessages).build(); + + // 4. Add the new user input to the conversation memory. + UserMessage userMessage = new UserMessage(request.userText(), request.media()); + this.getChatMemoryStore().add(this.doGetConversationId(context), userMessage); + + return advisedRequest; + } + + @Override + public ChatResponse adviseResponse(ChatResponse chatResponse, Map context) { + + List assistantMessages = chatResponse.getResults().stream().map(g -> (Message) g.getOutput()).toList(); + + this.getChatMemoryStore().add(this.doGetConversationId(context), assistantMessages); + + return chatResponse; + } + + @Override + public Flux adviseResponse(Flux fluxChatResponse, Map context) { + + return new MessageAggregator().aggregate(fluxChatResponse, chatResponse -> { + List assistantMessages = chatResponse.getResults() + .stream() + .map(g -> (Message) g.getOutput()) + .toList(); + + this.getChatMemoryStore().add(this.doGetConversationId(context), assistantMessages); + }); + } + +} \ No newline at end of file diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/PromptChatMemoryAdvisor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/PromptChatMemoryAdvisor.java new file mode 100644 index 000000000..822071b05 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/PromptChatMemoryAdvisor.java @@ -0,0 +1,124 @@ +/* + * Copyright 2024-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; + +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.AdvisedRequest; +import org.springframework.ai.chat.memory.ChatMemory; +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.ChatResponse; +import org.springframework.ai.chat.model.MessageAggregator; + +/** + * Memory is retrieved added into the prompt's system text. + * + * @author Christian Tzolov + * @since 1.0.0 M1 + */ +public class PromptChatMemoryAdvisor extends AbstractChatMemoryAdvisor { + + private static final String DEFAULT_SYSTEM_TEXT_ADVISE = """ + + Use the conversation memory from the MEMORY section to provide accurate answers. + + --------------------- + MEMORY: + {memory} + --------------------- + + """; + + private final String systemTextAdvise; + + public PromptChatMemoryAdvisor(ChatMemory chatMemory) { + this(chatMemory, DEFAULT_SYSTEM_TEXT_ADVISE); + } + + public PromptChatMemoryAdvisor(ChatMemory chatMemory, String systemTextAdvise) { + super(chatMemory); + this.systemTextAdvise = systemTextAdvise; + } + + public PromptChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int chatHistoryWindowSize, + String systemTextAdvise) { + super(chatMemory, defaultConversationId, chatHistoryWindowSize); + this.systemTextAdvise = systemTextAdvise; + } + + @Override + public AdvisedRequest adviseRequest(AdvisedRequest request, Map context) { + + // 1. Advise system parameters. + List memoryMessages = this.getChatMemoryStore() + .get(this.doGetConversationId(context), this.doGetChatMemoryRetrieveSize(context)); + + String memory = (memoryMessages != null) ? memoryMessages.stream() + .filter(m -> m.getMessageType() != MessageType.SYSTEM) + .map(m -> m.getMessageType() + ":" + m.getContent()) + .collect(Collectors.joining(System.lineSeparator())) : ""; + + Map advisedSystemParams = new HashMap<>(request.systemParams()); + advisedSystemParams.put("memory", memory); + + // 2. Advise the system text. + String advisedSystemText = request.systemText() + System.lineSeparator() + this.systemTextAdvise; + + // 3. Create a new request with the advised system text and parameters. + AdvisedRequest advisedRequest = AdvisedRequest.from(request) + .withSystemText(advisedSystemText) + .withSystemParams(advisedSystemParams) + .build(); + + // 4. Add the new user input to the conversation memory. + UserMessage userMessage = new UserMessage(request.userText(), request.media()); + this.getChatMemoryStore().add(this.doGetConversationId(context), userMessage); + + return advisedRequest; + } + + @Override + public ChatResponse adviseResponse(ChatResponse chatResponse, Map context) { + + List assistantMessages = chatResponse.getResults().stream().map(g -> (Message) g.getOutput()).toList(); + + this.getChatMemoryStore().add(this.doGetConversationId(context), assistantMessages); + + return chatResponse; + } + + @Override + public Flux adviseResponse(Flux fluxChatResponse, Map context) { + + return new MessageAggregator().aggregate(fluxChatResponse, chatResponse -> { + List assistantMessages = chatResponse.getResults() + .stream() + .map(g -> (Message) g.getOutput()) + .toList(); + + this.getChatMemoryStore().add(this.doGetConversationId(context), assistantMessages); + }); + } + +} \ No newline at end of file diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/QuestionAnswerAdvisor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/QuestionAnswerAdvisor.java new file mode 100644 index 000000000..9db7d9461 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/QuestionAnswerAdvisor.java @@ -0,0 +1,98 @@ +/* + * Copyright 2024-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; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import org.springframework.ai.chat.client.AdvisedRequest; +import org.springframework.ai.chat.client.RequestResponseAdvisor; +import org.springframework.ai.document.Document; +import org.springframework.ai.model.Content; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.util.Assert; + +/** + * Context for the question is retrieved from a Vector Store and added to the prompt's + * user text. + * + * @author Christian Tzolov + * @since 1.0.0 M1 + */ +public class QuestionAnswerAdvisor implements RequestResponseAdvisor { + + private static final String DEFAULT_USER_TEXT_ADVISE = """ + Context information is below. + --------------------- + {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 final VectorStore vectorStore; + + private final String userTextAdvise; + + private final SearchRequest searchRequest; + + public QuestionAnswerAdvisor(VectorStore vectorStore, SearchRequest searchRequest) { + this(vectorStore, searchRequest, DEFAULT_USER_TEXT_ADVISE); + } + + public QuestionAnswerAdvisor(VectorStore vectorStore, SearchRequest searchRequest, String userTextAdvise) { + + 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; + } + + @Override + public AdvisedRequest adviseRequest(AdvisedRequest request, Map context) { + + // 1. Advise the system text. + String advisedUserText = request.userText() + System.lineSeparator() + this.userTextAdvise; + + // 2. Search for similar documents in the vector store. + List documents = vectorStore.similaritySearch(searchRequest.withQuery(request.userText())); + + // 3. Create the context from the documents. + String documentContext = documents.stream() + .map(Content::getContent) + .collect(Collectors.joining(System.lineSeparator())); + + // 4. Advise the user parameters. + Map advisedUserParams = new HashMap<>(request.userParams()); + advisedUserParams.put("context", documentContext); + + AdvisedRequest advisedRequest = AdvisedRequest.from(request) + .withUserText(advisedUserText) + .withUserParams(advisedUserParams) + .build(); + + return advisedRequest; + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/VectorStoreChatMemoryAdvisor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/VectorStoreChatMemoryAdvisor.java new file mode 100644 index 000000000..c1809e85e --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/VectorStoreChatMemoryAdvisor.java @@ -0,0 +1,145 @@ +/* + * Copyright 2024-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; + +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.AdvisedRequest; +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.ChatResponse; +import org.springframework.ai.chat.model.MessageAggregator; +import org.springframework.ai.document.Document; +import org.springframework.ai.model.Content; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; + +/** + * Memory is retrieved from a VectorStore added into the prompt's system text. + * + * @author Christian Tzolov + * @since 1.0.0 M1 + */ +public class VectorStoreChatMemoryAdvisor extends AbstractChatMemoryAdvisor { + + 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, String systemTextAdvise) { + super(vectorStore, defaultConversationId, chatHistoryWindowSize); + this.systemTextAdvise = systemTextAdvise; + } + + @Override + public AdvisedRequest adviseRequest(AdvisedRequest request, Map context) { + + String advisedSystemText = request.systemText() + System.lineSeparator() + this.systemTextAdvise; + + var searchRequest = SearchRequest.query(request.userText()) + .withTopK(this.doGetChatMemoryRetrieveSize(context)) + .withFilterExpression(DOCUMENT_METADATA_CONVERSATION_ID + "=='" + this.doGetConversationId(context) + "'"); + + List documents = this.getChatMemoryStore().similaritySearch(searchRequest); + + String longTermMemory = documents.stream() + .map(Content::getContent) + .collect(Collectors.joining(System.lineSeparator())); + + Map advisedSystemParams = new HashMap<>(request.systemParams()); + advisedSystemParams.put("long_term_memory", longTermMemory); + + AdvisedRequest advisedRequest = AdvisedRequest.from(request) + .withSystemText(advisedSystemText) + .withSystemParams(advisedSystemParams) + .build(); + + UserMessage userMessage = new UserMessage(request.userText(), request.media()); + this.getChatMemoryStore().write(toDocuments(List.of(userMessage), this.doGetConversationId(context))); + + return advisedRequest; + } + + @Override + public ChatResponse adviseResponse(ChatResponse chatResponse, Map context) { + + List assistantMessages = chatResponse.getResults().stream().map(g -> (Message) g.getOutput()).toList(); + + this.getChatMemoryStore().write(toDocuments(assistantMessages, this.doGetConversationId(context))); + + return chatResponse; + } + + @Override + public Flux adviseResponse(Flux fluxChatResponse, Map context) { + + return new MessageAggregator().aggregate(fluxChatResponse, chatResponse -> { + List assistantMessages = chatResponse.getResults() + .stream() + .map(g -> (Message) g.getOutput()) + .toList(); + + this.getChatMemoryStore().write(toDocuments(assistantMessages, this.doGetConversationId(context))); + }); + } + + private List toDocuments(List messages, String conversationId) { + + List 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()); + var doc = new Document(message.getContent(), metadata); + return doc; + }) + .toList(); + + return docs; + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/ChatMemoryChatServiceListener.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/ChatMemoryChatServiceListener.java index 934c3d43d..7e6a4a44f 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/ChatMemoryChatServiceListener.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/ChatMemoryChatServiceListener.java @@ -20,14 +20,19 @@ import java.util.List; import org.springframework.ai.chat.service.ChatServiceResponse; import org.springframework.ai.chat.service.ChatServiceListener; +import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor; +import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor; import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.messages.MessageType; import org.springframework.ai.chat.prompt.transformer.ChatServiceContext; import org.springframework.ai.chat.prompt.transformer.TransformerContentType; /** + * @deprecated Use the {@link MessageChatMemoryAdvisor} or {@link PromptChatMemoryAdvisor} + * instead. * @author Christian Tzolov */ +@Deprecated public class ChatMemoryChatServiceListener implements ChatServiceListener { private final ChatMemory chatHistory; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/InMemoryChatMemory.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/InMemoryChatMemory.java index 34b78963e..4e8578e2f 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/InMemoryChatMemory.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/InMemoryChatMemory.java @@ -24,7 +24,16 @@ import java.util.concurrent.ConcurrentHashMap; import org.springframework.ai.chat.messages.Message; /** + * The InMemoryChatMemory class is an implementation of the ChatMemory interface that + * represents an in-memory storage for chat conversation history. + * + * This class stores the conversation history in a ConcurrentHashMap, where the keys are + * the conversation IDs and the values are lists of messages representing the conversation + * history. + * + * @see ChatMemory * @author Christian Tzolov + * @since 1.0.0 M1 */ public class InMemoryChatMemory implements ChatMemory { diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/LastMaxTokenSizeContentTransformer.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/LastMaxTokenSizeContentTransformer.java index b2e1626e6..903f6e3ba 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/LastMaxTokenSizeContentTransformer.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/LastMaxTokenSizeContentTransformer.java @@ -20,19 +20,23 @@ import java.util.ArrayList; import java.util.List; import java.util.Set; +import org.springframework.ai.chat.client.advisor.LastMaxTokenSizeContentPurger; import org.springframework.ai.chat.prompt.transformer.AbstractPromptTransformer; import org.springframework.ai.chat.prompt.transformer.ChatServiceContext; import org.springframework.ai.model.Content; import org.springframework.ai.tokenizer.TokenCountEstimator; /** + * * Returns a new list of content (e.g list of messages of list of documents) that is a * subset of the input list of contents and complies with the max token size constraint. * * The token estimator is used to estimate the token count of the datum. * + * @deprecated Use the {@link LastMaxTokenSizeContentPurger} instead. * @author Christian Tzolov */ +@Deprecated public class LastMaxTokenSizeContentTransformer extends AbstractPromptTransformer { protected final TokenCountEstimator tokenCountEstimator; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/MessageChatMemoryAugmentor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/MessageChatMemoryAugmentor.java index e3b3ce847..4c69b7c89 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/MessageChatMemoryAugmentor.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/MessageChatMemoryAugmentor.java @@ -19,6 +19,7 @@ package org.springframework.ai.chat.memory; import java.util.ArrayList; import java.util.List; +import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor; import org.springframework.ai.chat.messages.AbstractMessage; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.Message; @@ -32,8 +33,10 @@ import org.springframework.ai.chat.prompt.transformer.PromptChange; import org.springframework.ai.chat.prompt.transformer.TransformerContentType; /** + * @deprecated Use the {@link MessageChatMemoryAdvisor} instead. * @author Christian Tzolov */ +@Deprecated public class MessageChatMemoryAugmentor extends AbstractPromptTransformer { @Override diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/SystemPromptChatMemoryAugmentor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/SystemPromptChatMemoryAugmentor.java index 2ed8306ee..3fc6ee0dd 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/SystemPromptChatMemoryAugmentor.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/SystemPromptChatMemoryAugmentor.java @@ -22,6 +22,7 @@ import java.util.List; import java.util.Set; import java.util.stream.Collectors; +import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor; import org.springframework.ai.chat.messages.AbstractMessage; import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.messages.MessageType; @@ -35,8 +36,10 @@ import org.springframework.ai.chat.prompt.transformer.TransformerContentType; import org.springframework.util.Assert; /** + * @deprecated Use the {@link PromptChatMemoryAdvisor} instead. * @author Christian Tzolov */ +@Deprecated public class SystemPromptChatMemoryAugmentor extends AbstractPromptTransformer { public static final String DEFAULT_HISTORY_PROMPT = """ diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/VectorStoreChatMemoryChatServiceListener.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/VectorStoreChatMemoryChatServiceListener.java index 0f5aa0100..ee5d52eb3 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/VectorStoreChatMemoryChatServiceListener.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/VectorStoreChatMemoryChatServiceListener.java @@ -22,6 +22,7 @@ import java.util.Map; import org.springframework.ai.chat.service.ChatServiceListener; import org.springframework.ai.chat.service.ChatServiceResponse; +import org.springframework.ai.chat.client.advisor.VectorStoreChatMemoryAdvisor; import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.messages.MessageType; import org.springframework.ai.chat.prompt.transformer.TransformerContentType; @@ -31,8 +32,10 @@ import org.springframework.ai.vectorstore.VectorStore; import org.springframework.util.CollectionUtils; /** + * @deprecated Use the {@link VectorStoreChatMemoryAdvisor} instead. * @author Christian Tzolov */ +@Deprecated public class VectorStoreChatMemoryChatServiceListener implements ChatServiceListener { private final VectorStore vectorStore; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/VectorStoreChatMemoryRetriever.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/VectorStoreChatMemoryRetriever.java index 7b476e704..3a4914190 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/VectorStoreChatMemoryRetriever.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/VectorStoreChatMemoryRetriever.java @@ -21,6 +21,7 @@ import java.util.List; import java.util.Map; import java.util.stream.Collectors; +import org.springframework.ai.chat.client.advisor.VectorStoreChatMemoryAdvisor; import org.springframework.ai.chat.messages.MessageType; import org.springframework.ai.chat.prompt.transformer.AbstractPromptTransformer; import org.springframework.ai.chat.prompt.transformer.ChatServiceContext; @@ -32,8 +33,10 @@ import org.springframework.ai.vectorstore.VectorStore; import org.springframework.util.CollectionUtils; /** + * @deprecated Use the {@link VectorStoreChatMemoryAdvisor} instead. * @author Christian Tzolov */ +@Deprecated public class VectorStoreChatMemoryRetriever extends AbstractPromptTransformer { private final VectorStore vectorStore; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/model/ChatModel.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/model/ChatModel.java index 3fe6fb48d..5f2687bca 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/model/ChatModel.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/model/ChatModel.java @@ -20,11 +20,13 @@ import org.springframework.ai.chat.prompt.Prompt; import java.util.Arrays; +import reactor.core.publisher.Flux; + import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.messages.UserMessage; import org.springframework.ai.model.Model; -public interface ChatModel extends Model { +public interface ChatModel extends Model, StreamingChatModel { default String call(String message) { Prompt prompt = new Prompt(new UserMessage(message)); @@ -43,4 +45,8 @@ public interface ChatModel extends Model { ChatOptions getDefaultOptions(); + default Flux stream(Prompt prompt) { + throw new UnsupportedOperationException("streaming is not supported"); + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/service/MessageAggregator.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/model/MessageAggregator.java similarity index 94% rename from spring-ai-core/src/main/java/org/springframework/ai/chat/service/MessageAggregator.java rename to spring-ai-core/src/main/java/org/springframework/ai/chat/model/MessageAggregator.java index 849fa0b05..b4f688f87 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/service/MessageAggregator.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/model/MessageAggregator.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.chat.service; +package org.springframework.ai.chat.model; import java.util.HashMap; import java.util.List; @@ -26,9 +26,6 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; -import org.springframework.ai.chat.model.ChatResponse; -import org.springframework.ai.chat.model.Generation; - /** * Helper that for streaming chat responses, aggregate the chat response messages into a * single AssistantMessage. Job is performed in parallel to the chat response processing. diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/service/StreamingPromptTransformingChatService.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/service/StreamingPromptTransformingChatService.java index 6ddd29090..49381eb75 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/service/StreamingPromptTransformingChatService.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/service/StreamingPromptTransformingChatService.java @@ -23,6 +23,7 @@ import org.springframework.ai.chat.prompt.transformer.ChatServiceContext; import reactor.core.publisher.Flux; import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.MessageAggregator; import org.springframework.ai.chat.model.StreamingChatModel; import org.springframework.ai.chat.prompt.transformer.PromptTransformer; diff --git a/spring-ai-core/src/test/java/org/springframework/ai/chat/client/ChatClientAdvisorTests.java b/spring-ai-core/src/test/java/org/springframework/ai/chat/client/ChatClientAdvisorTests.java new file mode 100644 index 000000000..432029fb8 --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/chat/client/ChatClientAdvisorTests.java @@ -0,0 +1,272 @@ +/* + * Copyright 2024-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; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +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 reactor.core.publisher.Flux; + +import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor; +import org.springframework.ai.chat.memory.ChatMemory; +import org.springframework.ai.chat.memory.InMemoryChatMemory; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.MessageType; +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 static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.when; + +/** + * @author Christian Tzolov + */ +@ExtendWith(MockitoExtension.class) +public class ChatClientAdvisorTests { + + @Mock + ChatModel chatModel; + + @Captor + ArgumentCaptor promptCaptor; + + private String join(Flux fluxContent) { + return fluxContent.collectList().block().stream().collect(Collectors.joining()); + } + + @Test + public void promptChatMemory() { + + when(chatModel.call(promptCaptor.capture())) + .thenReturn(new ChatResponse(List.of(new Generation("Hello John")))) + .thenReturn(new ChatResponse(List.of(new Generation("Your name is John")))); + + ChatMemory chatMemory = new InMemoryChatMemory(); + + var chatClient = ChatClient.builder(chatModel) + .defaultSystem("Default system text.") + .defaultAdvisors(new PromptChatMemoryAdvisor(chatMemory)) + .build(); + + var content = chatClient.prompt() + .user("my name is John") + .call().content(); + + assertThat(content).isEqualTo("Hello John"); + + Message systemMessage = promptCaptor.getValue().getInstructions().get(0); + assertThat(systemMessage.getContent()).isEqualToIgnoringWhitespace(""" + Default system text. + + Use the conversation memory from the MEMORY section to provide accurate answers. + + --------------------- + MEMORY: + --------------------- + """); + assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM); + + Message userMessage = promptCaptor.getValue().getInstructions().get(1); + assertThat(userMessage.getContent()).isEqualToIgnoringWhitespace("my name is John"); + + content = chatClient.prompt() + .user("What is my name?") + .call().content(); + + assertThat(content).isEqualTo("Your name is John"); + + systemMessage = promptCaptor.getValue().getInstructions().get(0); + assertThat(systemMessage.getContent()).isEqualToIgnoringWhitespace(""" + Default system text. + + Use the conversation memory from the MEMORY section to provide accurate answers. + + --------------------- + MEMORY: + USER:my name is John + ASSISTANT:Hello John + --------------------- + """); + assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM); + + userMessage = promptCaptor.getValue().getInstructions().get(1); + assertThat(userMessage.getContent()).isEqualToIgnoringWhitespace("What is my name?"); + } + + @Test + public void streamingPromptChatMemory() { + + when(chatModel.stream(promptCaptor.capture())) + .thenReturn( + Flux.generate(() -> new ChatResponse(List.of(new Generation("Hello John"))), (state, sink) -> { + sink.next(state); + sink.complete(); + return state; + })) + .thenReturn( + Flux.generate(() -> new ChatResponse(List.of(new Generation("Your name is John"))), + (state, sink) -> { + sink.next(state); + sink.complete(); + return state; + })); + + ChatMemory chatMemory = new InMemoryChatMemory(); + + var chatClient = ChatClient.builder(chatModel) + .defaultSystem("Default system text.") + .defaultAdvisors(new PromptChatMemoryAdvisor(chatMemory)) + .build(); + + var content = join(chatClient.prompt() + .user("my name is John") + .stream().content()); + + assertThat(content).isEqualTo("Hello John"); + + Message systemMessage = promptCaptor.getValue().getInstructions().get(0); + assertThat(systemMessage.getContent()).isEqualToIgnoringWhitespace(""" + Default system text. + + Use the conversation memory from the MEMORY section to provide accurate answers. + + --------------------- + MEMORY: + --------------------- + """); + assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM); + + Message userMessage = promptCaptor.getValue().getInstructions().get(1); + assertThat(userMessage.getContent()).isEqualToIgnoringWhitespace("my name is John"); + + content = join(chatClient.prompt() + .user("What is my name?") + .stream().content()); + + assertThat(content).isEqualTo("Your name is John"); + + systemMessage = promptCaptor.getValue().getInstructions().get(0); + assertThat(systemMessage.getContent()).isEqualToIgnoringWhitespace(""" + Default system text. + + Use the conversation memory from the MEMORY section to provide accurate answers. + + --------------------- + MEMORY: + USER:my name is John + ASSISTANT:Hello John + --------------------- + """); + assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM); + + userMessage = promptCaptor.getValue().getInstructions().get(1); + assertThat(userMessage.getContent()).isEqualToIgnoringWhitespace("What is my name?"); + } + + public static class MockAdvisor implements RequestResponseAdvisor { + + public AdvisedRequest advisedRequest; + + public Map advisedRequestContext; + + public Map chatResponseContext; + + public ChatResponse chatResponse; + + public Map fluxChatResponseContext; + + public Flux fluxChatResponse; + + @Override + public AdvisedRequest adviseRequest(AdvisedRequest request, Map context) { + advisedRequest = request; + advisedRequestContext = context; + + context.put("adviseRequest", "adviseRequest"); + + return request; + } + + @Override + public ChatResponse adviseResponse(ChatResponse response, Map context) { + chatResponse = response; + chatResponseContext = context; + + context.put("adviseResponse", "adviseResponse"); + return response; + } + + @Override + public Flux adviseResponse(Flux fluxResponse, Map context) { + fluxChatResponse = fluxResponse; + fluxChatResponseContext = context; + + context.put("fluxAdviseResponse", "fluxAdviseResponse"); + + return fluxResponse; + } + + }; + + @Test + public void advisors() { + + var mockAdvisor = new MockAdvisor(); + + when(chatModel.call(promptCaptor.capture())).thenReturn(new ChatResponse(List.of(new Generation("Hello John")))) + .thenReturn(new ChatResponse(List.of(new Generation("Your name is John")))); + + when(chatModel.call(promptCaptor.capture())).thenReturn(new ChatResponse(List.of(new Generation("Hello John")))) + .thenReturn(new ChatResponse(List.of(new Generation("Your name is John")))); + + var chatClient = ChatClient.builder(chatModel) + .defaultSystem("Default system text.") + .defaultAdvisors(mockAdvisor) + .build(); + + var content = chatClient.prompt() + .user("my name is John") + .advisors(a -> a.param("key1", "value1").params(Map.of("key2", "value2"))) + .call() + .content(); + + assertThat(content).isEqualTo("Hello John"); + + assertThat(mockAdvisor.advisedRequestContext).containsEntry("key1", "value1") + .containsEntry("key2", "value2") + .containsEntry("adviseRequest", "adviseRequest"); + assertThat(mockAdvisor.advisedRequest.advisorParams()).containsEntry("key1", "value1") + .containsEntry("key2", "value2") + .doesNotContainKey("adviseRequest"); + + assertThat(mockAdvisor.chatResponseContext).containsEntry("key1", "value1") + .containsEntry("key2", "value2") + .containsEntry("adviseRequest", "adviseRequest") + .containsEntry("adviseResponse", "adviseResponse"); + assertThat(mockAdvisor.chatResponse).isNotNull(); + } + +} diff --git a/spring-ai-core/src/test/java/org/springframework/ai/chat/client/ChatClientTest.java b/spring-ai-core/src/test/java/org/springframework/ai/chat/client/ChatClientTest.java index ac44c9dc7..3f05bfb1a 100644 --- a/spring-ai-core/src/test/java/org/springframework/ai/chat/client/ChatClientTest.java +++ b/spring-ai-core/src/test/java/org/springframework/ai/chat/client/ChatClientTest.java @@ -30,13 +30,12 @@ import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import reactor.core.publisher.Flux; -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.model.StreamingChatModel; 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.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.model.function.FunctionCallingOptions; import org.springframework.ai.model.function.FunctionCallingOptionsBuilder; @@ -53,12 +52,8 @@ import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) public class ChatClientTest { - public static interface MixChatModel extends ChatModel, StreamingChatModel { - - } - @Mock - MixChatModel chatModel; + ChatModel chatModel; @Captor ArgumentCaptor promptCaptor;