Skip to content

Commit 55cbf15

Browse files
committed
Introduce first-class chat memory support
- ChatMemory will become a generic interface to implement different memory management strategies. It’s been moved from the “”spring-ai-client-chat” package to “spring-ai-model” package while retaining the same package, so it’s transparent to users. - A MessageWindowChatMemory has been introduced to provide support for a chat memory that keeps at most N messages in the memory. - A MessageWindowProcessingPolicy API has been introduced to customise the processing policy for the message window. A default implementation is provided out-of-the-box. - A ChatMemoryRepository interface has been introduced to support different storage strategies for the chat memory. It’s meant to be used as part of a ChatMemory implementation. This is different than before, where the storage-specific implementation was directly tied to the ChatMemory. This design is familiar to Spring users since it’s used already in the ecosystem. The goal was to use a programming model similar to Spring Session and Spring Data. - The JdbcChatMemory has been supersed by JdbcChatMemoryRepository. - The ChatClient now supports memory as a first-class citizen, superseding the need for an Advisor to manage the chat memory. It also simplifies providing a conversationId. This feature lays the foundation for including the intermediate messages in tool calling in the memory as well. - All the changes introduced in this PR are backword-compatible. Signed-off-by: Thomas Vitale <[email protected]>
1 parent c0bc623 commit 55cbf15

File tree

38 files changed

+1754
-70
lines changed

38 files changed

+1754
-70
lines changed

auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/src/main/java/org/springframework/ai/model/chat/memory/jdbc/autoconfigure/JdbcChatMemoryAutoConfiguration.java

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323

2424
import org.springframework.ai.chat.memory.jdbc.JdbcChatMemory;
2525
import org.springframework.ai.chat.memory.jdbc.JdbcChatMemoryConfig;
26+
import org.springframework.ai.chat.memory.jdbc.JdbcChatMemoryRepository;
2627
import org.springframework.boot.autoconfigure.AutoConfiguration;
2728
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
2829
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
@@ -35,6 +36,7 @@
3536

3637
/**
3738
* @author Jonathan Leijendekker
39+
* @author Thomas Vitale
3840
* @since 1.0.0
3941
*/
4042
@AutoConfiguration(after = JdbcTemplateAutoConfiguration.class)
@@ -46,19 +48,30 @@ public class JdbcChatMemoryAutoConfiguration {
4648

4749
@Bean
4850
@ConditionalOnMissingBean
49-
public JdbcChatMemory chatMemory(JdbcTemplate jdbcTemplate) {
51+
JdbcChatMemoryRepository jdbcChatMemoryRepository(JdbcTemplate jdbcTemplate) {
5052
var config = JdbcChatMemoryConfig.builder().jdbcTemplate(jdbcTemplate).build();
53+
return JdbcChatMemoryRepository.create(config);
54+
}
5155

56+
/**
57+
* @deprecated in favor of providing ChatClient directly a
58+
* {@link org.springframework.ai.chat.memory.MessageWindowChatMemory} with a
59+
* {@link JdbcChatMemoryRepository} instance.
60+
*/
61+
@Bean
62+
@ConditionalOnMissingBean
63+
@Deprecated
64+
JdbcChatMemory chatMemory(JdbcTemplate jdbcTemplate) {
65+
var config = JdbcChatMemoryConfig.builder().jdbcTemplate(jdbcTemplate).build();
5266
return JdbcChatMemory.create(config);
5367
}
5468

5569
@Bean
5670
@ConditionalOnMissingBean
5771
@ConditionalOnProperty(value = "spring.ai.chat.memory.jdbc.initialize-schema", havingValue = "true",
5872
matchIfMissing = true)
59-
public DataSourceScriptDatabaseInitializer jdbcChatMemoryScriptDatabaseInitializer(DataSource dataSource) {
73+
DataSourceScriptDatabaseInitializer jdbcChatMemoryScriptDatabaseInitializer(DataSource dataSource) {
6074
logger.debug("Initializing JdbcChatMemory schema");
61-
6275
return new JdbcChatMemoryDataSourceScriptDatabaseInitializer(dataSource);
6376
}
6477

auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/src/test/java/org/springframework/ai/model/chat/memory/jdbc/autoconfigure/JdbcChatMemoryAutoConfigurationIT.java

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import org.testcontainers.utility.DockerImageName;
2727

2828
import org.springframework.ai.chat.memory.jdbc.JdbcChatMemory;
29+
import org.springframework.ai.chat.memory.jdbc.JdbcChatMemoryRepository;
2930
import org.springframework.ai.chat.messages.AssistantMessage;
3031
import org.springframework.ai.chat.messages.Message;
3132
import org.springframework.ai.chat.messages.UserMessage;
@@ -38,6 +39,7 @@
3839

3940
/**
4041
* @author Jonathan Leijendekker
42+
* @author Thomas Vitale
4143
*/
4244
@Testcontainers
4345
class JdbcChatMemoryAutoConfigurationIT {
@@ -96,4 +98,30 @@ void addGetAndClear_shouldAllExecute() {
9698
});
9799
}
98100

101+
@Test
102+
void useAutoConfiguredJdbcChatMemoryRepository() {
103+
this.contextRunner.withPropertyValues("spring.ai.chat.memory.jdbc.initialize-schema=true").run(context -> {
104+
var chatMemoryRepository = context.getBean(JdbcChatMemoryRepository.class);
105+
var conversationId = UUID.randomUUID().toString();
106+
var userMessage = new UserMessage("Message from the user");
107+
108+
chatMemoryRepository.save(conversationId, List.of(userMessage));
109+
110+
assertThat(chatMemoryRepository.findById(conversationId)).hasSize(1);
111+
assertThat(chatMemoryRepository.findById(conversationId)).isEqualTo(List.of(userMessage));
112+
113+
chatMemoryRepository.deleteById(conversationId);
114+
115+
assertThat(chatMemoryRepository.findById(conversationId)).isEmpty();
116+
117+
var multipleMessages = List.<Message>of(new UserMessage("Message from the user 1"),
118+
new AssistantMessage("Message from the assistant 1"));
119+
120+
chatMemoryRepository.save(conversationId, multipleMessages);
121+
122+
assertThat(chatMemoryRepository.findById(conversationId)).hasSize(multipleMessages.size());
123+
assertThat(chatMemoryRepository.findById(conversationId)).isEqualTo(multipleMessages);
124+
});
125+
}
126+
99127
}

memory/spring-ai-model-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/jdbc/JdbcChatMemory.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,11 @@
3838
*
3939
* @author Jonathan Leijendekker
4040
* @since 1.0.0
41+
* @deprecated in favor of providing ChatClient directly a
42+
* {@link org.springframework.ai.chat.memory.MessageWindowChatMemory} with a
43+
* {@link JdbcChatMemoryRepository} instance.
4144
*/
45+
@Deprecated
4246
public class JdbcChatMemory implements ChatMemory {
4347

4448
private static final String QUERY_ADD = """

memory/spring-ai-model-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/jdbc/JdbcChatMemoryConfig.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
import org.springframework.util.Assert;
2121

2222
/**
23-
* Configuration for {@link JdbcChatMemory}.
23+
* Configuration for {@link JdbcChatMemoryRepository}.
2424
*
2525
* @author Jonathan Leijendekker
2626
* @since 1.0.0
Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
/*
2+
* Copyright 2023-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.chat.memory.jdbc;
18+
19+
import org.springframework.ai.chat.memory.ChatMemoryRepository;
20+
import org.springframework.ai.chat.messages.*;
21+
import org.springframework.jdbc.core.BatchPreparedStatementSetter;
22+
import org.springframework.jdbc.core.JdbcTemplate;
23+
import org.springframework.jdbc.core.RowMapper;
24+
import org.springframework.lang.Nullable;
25+
import org.springframework.util.Assert;
26+
27+
import java.sql.PreparedStatement;
28+
import java.sql.ResultSet;
29+
import java.sql.SQLException;
30+
import java.util.List;
31+
32+
/**
33+
* An implementation of {@link ChatMemoryRepository} for JDBC.
34+
*
35+
* @author Jonathan Leijendekker
36+
* @author Thomas Vitale
37+
* @since 1.0.0
38+
*/
39+
public class JdbcChatMemoryRepository implements ChatMemoryRepository {
40+
41+
private static final String QUERY_ADD = """
42+
INSERT INTO ai_chat_memory (conversation_id, content, type) VALUES (?, ?, ?)""";
43+
44+
private static final String QUERY_GET = """
45+
SELECT content, type FROM ai_chat_memory WHERE conversation_id = ? ORDER BY "timestamp" DESC""";
46+
47+
private static final String QUERY_CLEAR = "DELETE FROM ai_chat_memory WHERE conversation_id = ?";
48+
49+
private final JdbcTemplate jdbcTemplate;
50+
51+
private JdbcChatMemoryRepository(JdbcChatMemoryConfig config) {
52+
Assert.notNull(config, "config cannot be null");
53+
this.jdbcTemplate = config.getJdbcTemplate();
54+
}
55+
56+
public static JdbcChatMemoryRepository create(JdbcChatMemoryConfig config) {
57+
return new JdbcChatMemoryRepository(config);
58+
}
59+
60+
@Override
61+
public List<Message> findById(String conversationId) {
62+
Assert.hasText(conversationId, "conversationId cannot be null or empty");
63+
return this.jdbcTemplate.query(QUERY_GET, new MessageRowMapper(), conversationId);
64+
}
65+
66+
@Override
67+
public void save(String conversationId, List<Message> messages) {
68+
Assert.hasText(conversationId, "conversationId cannot be null or empty");
69+
Assert.notNull(messages, "messages cannot be null");
70+
Assert.noNullElements(messages, "messages cannot contain null elements");
71+
this.jdbcTemplate.batchUpdate(QUERY_ADD, new AddBatchPreparedStatement(conversationId, messages));
72+
}
73+
74+
@Override
75+
public void deleteById(String conversationId) {
76+
Assert.hasText(conversationId, "conversationId cannot be null or empty");
77+
this.jdbcTemplate.update(QUERY_CLEAR, conversationId);
78+
}
79+
80+
private record AddBatchPreparedStatement(String conversationId,
81+
List<Message> messages) implements BatchPreparedStatementSetter {
82+
@Override
83+
public void setValues(PreparedStatement ps, int i) throws SQLException {
84+
var message = this.messages.get(i);
85+
86+
ps.setString(1, this.conversationId);
87+
ps.setString(2, message.getText());
88+
ps.setString(3, message.getMessageType().name());
89+
}
90+
91+
@Override
92+
public int getBatchSize() {
93+
return this.messages.size();
94+
}
95+
}
96+
97+
private static class MessageRowMapper implements RowMapper<Message> {
98+
99+
@Override
100+
@Nullable
101+
public Message mapRow(ResultSet rs, int i) throws SQLException {
102+
var content = rs.getString(1);
103+
var type = MessageType.valueOf(rs.getString(2));
104+
105+
return switch (type) {
106+
case USER -> new UserMessage(content);
107+
case ASSISTANT -> new AssistantMessage(content);
108+
case SYSTEM -> new SystemMessage(content);
109+
case TOOL -> null;
110+
};
111+
}
112+
113+
}
114+
115+
}
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
*
2828
* @author Jonathan Leijendekker
2929
*/
30-
class JdbcChatMemoryRuntimeHints implements RuntimeHintsRegistrar {
30+
class JdbcChatMemoryRepositoryRuntimeHints implements RuntimeHintsRegistrar {
3131

3232
@Override
3333
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
/*
2+
* Copyright 2023-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
@NonNullApi
18+
@NonNullFields
19+
package org.springframework.ai.chat.memory.jdbc;
20+
21+
import org.springframework.lang.NonNullApi;
22+
import org.springframework.lang.NonNullFields;
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
11
org.springframework.aot.hint.RuntimeHintsRegistrar=\
2-
org.springframework.ai.chat.memory.jdbc.aot.hint.JdbcChatMemoryRuntimeHints
2+
org.springframework.ai.chat.memory.jdbc.aot.hint.JdbcChatMemoryRepositoryRuntimeHints

0 commit comments

Comments
 (0)