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

Polish ChatClientAdvisorTests

Browse Source
This commit is contained in:
Mark Pollack
2025-05-03 14:26:17 -04:00
parent 846165d47f
commit 2cbfb22db6
1 changed files with 28 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/ChatClientAdvisorTests.java
View File

@@ -37,13 +37,16 @@ import org.springframework.ai.chat.metadata.ChatResponseMetadata;
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.MessageAggregator;
import org.springframework.ai.chat.prompt.Prompt;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.BDDMockito.given;
/**
* Tests for the ChatClient with a focus on verifying the handling of conversation memory
* and the integration of PromptChatMemoryAdvisor to ensure accurate responses based on
* previous interactions.
*
* @author Christian Tzolov
* @author Alexandros Pappas
*/
@@ -63,32 +66,33 @@ public class ChatClientAdvisorTests {
@Test
public void promptChatMemory() {
var builder = ChatResponseMetadata.builder()
.id("124")
.usage(new MessageAggregator.DefaultUsage(1, 2, 3))
.model("gpt4o")
.keyValue("created", 0L)
.keyValue("system-fingerprint", "john doe");
ChatResponseMetadata chatResponseMetadata = builder.build();
// Create a ChatResponseMetadata instance with default values
ChatResponseMetadata chatResponseMetadata = ChatResponseMetadata.builder().build();
// Mock the chatModel to return predefined ChatResponse objects when called
given(this.chatModel.call(this.promptCaptor.capture()))
.willReturn(
new ChatResponse(List.of(new Generation(new AssistantMessage("Hello John"))), chatResponseMetadata))
.willReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("Your name is John"))),
chatResponseMetadata));
// Initialize an in-memory chat memory to store conversation history
ChatMemory chatMemory = new InMemoryChatMemory();
// Build a ChatClient with default system text and a memory advisor
var chatClient = ChatClient.builder(this.chatModel)
.defaultSystem("Default system text.")
.defaultAdvisors(new PromptChatMemoryAdvisor(chatMemory))
.build();
// Simulate a user prompt and verify the response
ChatResponse chatResponse = chatClient.prompt().user("my name is John").call().chatResponse();
// Assert that the response content matches the expected output
String content = chatResponse.getResult().getOutput().getText();
assertThat(content).isEqualTo("Hello John");
// Capture and verify the system message instructions
Message systemMessage = this.promptCaptor.getValue().getInstructions().get(0);
assertThat(systemMessage.getText()).isEqualToIgnoringWhitespace("""
Default system text.
@@ -101,13 +105,17 @@ public class ChatClientAdvisorTests {
""");
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
// Capture and verify the user message instructions
Message userMessage = this.promptCaptor.getValue().getInstructions().get(1);
assertThat(userMessage.getText()).isEqualToIgnoringWhitespace("my name is John");
// Simulate another user prompt and verify the response
content = chatClient.prompt().user("What is my name?").call().content();
// Assert that the response content matches the expected output
assertThat(content).isEqualTo("Your name is John");
// Capture and verify the updated system message instructions
systemMessage = this.promptCaptor.getValue().getInstructions().get(0);
assertThat(systemMessage.getText()).isEqualToIgnoringWhitespace("""
Default system text.
@@ -122,6 +130,7 @@ public class ChatClientAdvisorTests {
""");
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
// Capture and verify the updated user message instructions
userMessage = this.promptCaptor.getValue().getInstructions().get(1);
assertThat(userMessage.getText()).isEqualToIgnoringWhitespace("What is my name?");
}
@@ -129,6 +138,7 @@ public class ChatClientAdvisorTests {
@Test
public void streamingPromptChatMemory() {
// Mock the chatModel to stream predefined ChatResponse objects
given(this.chatModel.stream(this.promptCaptor.capture())).willReturn(Flux.generate(
() -> new ChatResponse(List.of(new Generation(new AssistantMessage("Hello John")))), (state, sink) -> {
sink.next(state);
@@ -143,17 +153,22 @@ public class ChatClientAdvisorTests {
return state;
}));
// Initialize an in-memory chat memory to store conversation history
ChatMemory chatMemory = new InMemoryChatMemory();
// Build a ChatClient with default system text and a memory advisor
var chatClient = ChatClient.builder(this.chatModel)
.defaultSystem("Default system text.")
.defaultAdvisors(new PromptChatMemoryAdvisor(chatMemory))
.build();
// Simulate a streaming user prompt and verify the response
var content = join(chatClient.prompt().user("my name is John").stream().content());
// Assert that the streamed content matches the expected output
assertThat(content).isEqualTo("Hello John");
// Capture and verify the system message instructions
Message systemMessage = this.promptCaptor.getValue().getInstructions().get(0);
assertThat(systemMessage.getText()).isEqualToIgnoringWhitespace("""
Default system text.
@@ -166,13 +181,17 @@ public class ChatClientAdvisorTests {
""");
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
// Capture and verify the user message instructions
Message userMessage = this.promptCaptor.getValue().getInstructions().get(1);
assertThat(userMessage.getText()).isEqualToIgnoringWhitespace("my name is John");
// Simulate another streaming user prompt and verify the response
content = join(chatClient.prompt().user("What is my name?").stream().content());
// Assert that the streamed content matches the expected output
assertThat(content).isEqualTo("Your name is John");
// Capture and verify the updated system message instructions
systemMessage = this.promptCaptor.getValue().getInstructions().get(0);
assertThat(systemMessage.getText()).isEqualToIgnoringWhitespace("""
Default system text.
@@ -187,6 +206,7 @@ public class ChatClientAdvisorTests {
""");
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
// Capture and verify the updated user message instructions
userMessage = this.promptCaptor.getValue().getInstructions().get(1);
assertThat(userMessage.getText()).isEqualToIgnoringWhitespace("What is my name?");
}