Skip to content

Commit 772be8b

Browse files
ThomasVitaleThomasVitale
committed
Chat Memory Enhancements
* 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 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. * A ChatMemory bean is auto-configured for you whenever using one of the Spring AI Model starters. By default, it uses the MessageWindowChatMemory implementation and stores the conversation history in memory. If a different repository is already configured (e.g., Cassandra, JDBC, or Neo4j), Spring AI will use that instead. * First-class documentation has been introduced to describe the ChatMemory API and related features. * All the changes introduced in this PR are backward-compatible. Signed-off-by: Thomas Vitale <[email protected]>
1 parent 2bc29da commit 772be8b

File tree

60 files changed

+2107
-69
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

60 files changed

+2107
-69
lines changed

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-cassandra/pom.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,12 @@
3030
<version>${project.parent.version}</version>
3131
</dependency>
3232

33+
<dependency>
34+
<groupId>org.springframework.ai</groupId>
35+
<artifactId>spring-ai-autoconfigure-model-chat-memory</artifactId>
36+
<version>${project.parent.version}</version>
37+
</dependency>
38+
3339
<!-- Boot dependencies -->
3440
<dependency>
3541
<groupId>org.springframework.boot</groupId>

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-cassandra/src/main/java/org/springframework/ai/model/chat/memory/cassandra/autoconfigure/CassandraChatMemoryAutoConfiguration.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemory;
2222
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryConfig;
23+
import org.springframework.ai.model.chat.memory.autoconfigure.ChatMemoryAutoConfiguration;
2324
import org.springframework.boot.autoconfigure.AutoConfiguration;
2425
import org.springframework.boot.autoconfigure.cassandra.CassandraAutoConfiguration;
2526
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -34,7 +35,7 @@
3435
* @author Jihoon Kim
3536
* @since 1.0.0
3637
*/
37-
@AutoConfiguration(after = CassandraAutoConfiguration.class)
38+
@AutoConfiguration(after = CassandraAutoConfiguration.class, before = ChatMemoryAutoConfiguration.class)
3839
@ConditionalOnClass({ CassandraChatMemory.class, CqlSession.class })
3940
@EnableConfigurationProperties(CassandraChatMemoryProperties.class)
4041
public class CassandraChatMemoryAutoConfiguration {

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/pom.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,12 @@
2929
<version>${project.parent.version}</version>
3030
</dependency>
3131

32+
<dependency>
33+
<groupId>org.springframework.ai</groupId>
34+
<artifactId>spring-ai-autoconfigure-model-chat-memory</artifactId>
35+
<version>${project.parent.version}</version>
36+
</dependency>
37+
3238
<!-- Boot dependencies -->
3339
<dependency>
3440
<groupId>org.springframework.boot</groupId>

‎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: 59 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -21,45 +21,93 @@
2121
import org.slf4j.Logger;
2222
import org.slf4j.LoggerFactory;
2323

24+
import org.springframework.ai.chat.memory.ChatMemory;
25+
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
2426
import org.springframework.ai.chat.memory.jdbc.JdbcChatMemory;
2527
import org.springframework.ai.chat.memory.jdbc.JdbcChatMemoryConfig;
28+
import org.springframework.ai.chat.memory.jdbc.JdbcChatMemoryRepository;
29+
import org.springframework.ai.model.chat.memory.autoconfigure.ChatMemoryAutoConfiguration;
2630
import org.springframework.boot.autoconfigure.AutoConfiguration;
2731
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
2832
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
29-
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
33+
import org.springframework.boot.autoconfigure.condition.ConditionMessage;
34+
import org.springframework.boot.autoconfigure.condition.ConditionOutcome;
35+
import org.springframework.boot.autoconfigure.condition.SpringBootCondition;
3036
import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration;
3137
import org.springframework.boot.context.properties.EnableConfigurationProperties;
32-
import org.springframework.boot.jdbc.init.DataSourceScriptDatabaseInitializer;
3338
import org.springframework.context.annotation.Bean;
39+
import org.springframework.context.annotation.ConditionContext;
40+
import org.springframework.context.annotation.Conditional;
41+
import org.springframework.core.type.AnnotatedTypeMetadata;
3442
import org.springframework.jdbc.core.JdbcTemplate;
3543

3644
/**
3745
* @author Jonathan Leijendekker
46+
* @author Thomas Vitale
3847
* @since 1.0.0
3948
*/
40-
@AutoConfiguration(after = JdbcTemplateAutoConfiguration.class)
41-
@ConditionalOnClass({ JdbcChatMemory.class, DataSource.class, JdbcTemplate.class })
49+
@AutoConfiguration(after = JdbcTemplateAutoConfiguration.class, before = ChatMemoryAutoConfiguration.class)
50+
@ConditionalOnClass({ JdbcChatMemoryRepository.class, DataSource.class, JdbcTemplate.class })
4251
@EnableConfigurationProperties(JdbcChatMemoryProperties.class)
4352
public class JdbcChatMemoryAutoConfiguration {
4453

4554
private static final Logger logger = LoggerFactory.getLogger(JdbcChatMemoryAutoConfiguration.class);
4655

4756
@Bean
4857
@ConditionalOnMissingBean
49-
public JdbcChatMemory chatMemory(JdbcTemplate jdbcTemplate) {
50-
var config = JdbcChatMemoryConfig.builder().jdbcTemplate(jdbcTemplate).build();
58+
JdbcChatMemoryRepository chatMemoryRepository(JdbcTemplate jdbcTemplate) {
59+
return JdbcChatMemoryRepository.builder().jdbcTemplate(jdbcTemplate).build();
60+
}
5161

62+
/**
63+
* @deprecated in favor of building a {@link MessageWindowChatMemory} (or other
64+
* {@link ChatMemory} implementations) with a {@link JdbcChatMemoryRepository}
65+
* instance.
66+
*/
67+
@Bean
68+
@ConditionalOnMissingBean
69+
@Deprecated
70+
JdbcChatMemory chatMemory(JdbcTemplate jdbcTemplate) {
71+
var config = JdbcChatMemoryConfig.builder().jdbcTemplate(jdbcTemplate).build();
5272
return JdbcChatMemory.create(config);
5373
}
5474

5575
@Bean
5676
@ConditionalOnMissingBean
57-
@ConditionalOnProperty(value = "spring.ai.chat.memory.jdbc.initialize-schema", havingValue = "true",
58-
matchIfMissing = true)
59-
public DataSourceScriptDatabaseInitializer jdbcChatMemoryScriptDatabaseInitializer(DataSource dataSource) {
60-
logger.debug("Initializing JdbcChatMemory schema");
61-
77+
@Conditional(OnSchemaInitializationEnabledCondition.class)
78+
JdbcChatMemoryDataSourceScriptDatabaseInitializer jdbcChatMemoryScriptDatabaseInitializer(DataSource dataSource) {
79+
logger.debug("Initializing schema for JdbcChatMemoryRepository");
6280
return new JdbcChatMemoryDataSourceScriptDatabaseInitializer(dataSource);
6381
}
6482

83+
/**
84+
* Condition to check if the schema initialization is enabled, supporting both
85+
* deprecated and new property.
86+
*
87+
* @deprecated to be removed in 1.0.0-RC1.
88+
*/
89+
@Deprecated
90+
static class OnSchemaInitializationEnabledCondition extends SpringBootCondition {
91+
92+
@Override
93+
public ConditionOutcome getMatchOutcome(ConditionContext context, AnnotatedTypeMetadata metadata) {
94+
Boolean initializeSchemaEnabled = context.getEnvironment()
95+
.getProperty("spring.ai.chat.memory.jdbc.initialize-schema", Boolean.class);
96+
97+
if (initializeSchemaEnabled != null) {
98+
return new ConditionOutcome(initializeSchemaEnabled,
99+
ConditionMessage.forCondition("Enable JDBC Chat Memory Schema Initialization")
100+
.because("spring.ai.chat.memory.jdbc.initialize-schema is " + initializeSchemaEnabled));
101+
}
102+
103+
initializeSchemaEnabled = context.getEnvironment()
104+
.getProperty(JdbcChatMemoryProperties.CONFIG_PREFIX + ".initialize-schema", Boolean.class, true);
105+
106+
return new ConditionOutcome(initializeSchemaEnabled, ConditionMessage
107+
.forCondition("Enable JDBC Chat Memory Schema Initialization")
108+
.because(JdbcChatMemoryProperties.CONFIG_PREFIX + ".initialize-schema is " + initializeSchemaEnabled));
109+
}
110+
111+
}
112+
65113
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,11 @@
2525
import org.springframework.boot.sql.init.DatabaseInitializationMode;
2626
import org.springframework.boot.sql.init.DatabaseInitializationSettings;
2727

28+
/**
29+
* Performs database initialization for the JDBC Chat Memory Repository.
30+
*
31+
* @since 1.0.0
32+
*/
2833
class JdbcChatMemoryDataSourceScriptDatabaseInitializer extends DataSourceScriptDatabaseInitializer {
2934

3035
private static final String SCHEMA_LOCATION = "classpath:org/springframework/ai/chat/memory/jdbc/schema-@@platform@@.sql";

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,17 @@
2020

2121
/**
2222
* @author Jonathan Leijendekker
23+
* @author Thomas Vitale
2324
* @since 1.0.0
2425
*/
2526
@ConfigurationProperties(JdbcChatMemoryProperties.CONFIG_PREFIX)
2627
public class JdbcChatMemoryProperties {
2728

28-
public static final String CONFIG_PREFIX = "spring.ai.chat.memory.jdbc";
29+
public static final String CONFIG_PREFIX = "spring.ai.chat.memory.repository.jdbc";
2930

31+
/**
32+
* Whether to initialize the schema on startup.
33+
*/
3034
private boolean initializeSchema = true;
3135

3236
public boolean isInitializeSchema() {

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/src/main/resources/META-INF/additional-spring-configuration-metadata.json‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
{
2+
"groups": [],
3+
"properties": [
4+
{
5+
"name": "spring.ai.chat.memory.jdbc.initialize-schema",
6+
"type": "java.lang.Boolean",
7+
"description": "Whether to initialize the schema on startup.",
8+
"deprecation": {
9+
"replacement": "spring.ai.chat.memory.repository.jdbc.initialize-schema"
10+
}
11+
}
12+
],
13+
"hints": []
14+
}

‎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: 81 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -16,58 +16,63 @@
1616

1717
package org.springframework.ai.model.chat.memory.jdbc.autoconfigure;
1818

19-
import java.util.List;
20-
import java.util.UUID;
21-
2219
import org.junit.jupiter.api.Test;
23-
import org.testcontainers.containers.PostgreSQLContainer;
24-
import org.testcontainers.junit.jupiter.Container;
25-
import org.testcontainers.junit.jupiter.Testcontainers;
26-
import org.testcontainers.utility.DockerImageName;
27-
20+
import org.springframework.ai.chat.memory.ChatMemory;
2821
import org.springframework.ai.chat.memory.jdbc.JdbcChatMemory;
22+
import org.springframework.ai.chat.memory.jdbc.JdbcChatMemoryRepository;
2923
import org.springframework.ai.chat.messages.AssistantMessage;
3024
import org.springframework.ai.chat.messages.Message;
3125
import org.springframework.ai.chat.messages.UserMessage;
26+
import org.springframework.ai.model.chat.memory.autoconfigure.ChatMemoryAutoConfiguration;
3227
import org.springframework.boot.autoconfigure.AutoConfigurations;
3328
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
3429
import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration;
3530
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
3631

32+
import java.util.List;
33+
import java.util.UUID;
34+
3735
import static org.assertj.core.api.Assertions.assertThat;
3836

3937
/**
4038
* @author Jonathan Leijendekker
39+
* @author Thomas Vitale
4140
*/
42-
@Testcontainers
4341
class JdbcChatMemoryAutoConfigurationIT {
4442

45-
static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse("postgres:17");
46-
47-
@Container
48-
@SuppressWarnings("resource")
49-
static PostgreSQLContainer<?> postgresContainer = new PostgreSQLContainer<>(DEFAULT_IMAGE_NAME)
50-
.withDatabaseName("chat_memory_auto_configuration_test")
51-
.withUsername("postgres")
52-
.withPassword("postgres");
53-
5443
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
5544
.withConfiguration(AutoConfigurations.of(JdbcChatMemoryAutoConfiguration.class,
5645
JdbcTemplateAutoConfiguration.class, DataSourceAutoConfiguration.class))
57-
.withPropertyValues(String.format("spring.datasource.url=%s", postgresContainer.getJdbcUrl()),
58-
String.format("spring.datasource.username=%s", postgresContainer.getUsername()),
59-
String.format("spring.datasource.password=%s", postgresContainer.getPassword()));
46+
.withPropertyValues("spring.datasource.url=jdbc:tc:postgresql:17:///");
6047

6148
@Test
6249
void jdbcChatMemoryScriptDatabaseInitializer_shouldBeLoaded() {
6350
this.contextRunner.withPropertyValues("spring.ai.chat.memory.jdbc.initialize-schema=true")
6451
.run(context -> assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isTrue());
52+
this.contextRunner.withPropertyValues("spring.ai.chat.memory.repository.jdbc.initialize-schema=true")
53+
.run(context -> assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isTrue());
6554
}
6655

6756
@Test
6857
void jdbcChatMemoryScriptDatabaseInitializer_shouldNotBeLoaded() {
6958
this.contextRunner.withPropertyValues("spring.ai.chat.memory.jdbc.initialize-schema=false")
7059
.run(context -> assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isFalse());
60+
this.contextRunner.withPropertyValues("spring.ai.chat.memory.repository.jdbc.initialize-schema=false")
61+
.run(context -> assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isFalse());
62+
}
63+
64+
@Test
65+
void initializeSchemaEnabledWithDeprecatedProperty() {
66+
this.contextRunner
67+
.withPropertyValues("spring.ai.chat.memory.jdbc.initialize-schema=true",
68+
"spring.ai.chat.memory.repository.jdbc.initialize-schema=false")
69+
.run(context -> assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isTrue());
70+
}
71+
72+
@Test
73+
void initializeSchemaEnabledWithNewProperty() {
74+
this.contextRunner.withPropertyValues("spring.ai.chat.memory.repository.jdbc.initialize-schema=true")
75+
.run(context -> assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isTrue());
7176
}
7277

7378
@Test
@@ -96,4 +101,59 @@ void addGetAndClear_shouldAllExecute() {
96101
});
97102
}
98103

104+
@Test
105+
void useAutoConfiguredJdbcChatMemoryRepository() {
106+
this.contextRunner.run(context -> {
107+
var chatMemoryRepository = context.getBean(JdbcChatMemoryRepository.class);
108+
var conversationId = UUID.randomUUID().toString();
109+
var userMessage = new UserMessage("Message from the user");
110+
111+
chatMemoryRepository.saveMessages(conversationId, List.of(userMessage));
112+
113+
assertThat(chatMemoryRepository.findMessages(conversationId)).hasSize(1);
114+
assertThat(chatMemoryRepository.findMessages(conversationId)).isEqualTo(List.of(userMessage));
115+
116+
chatMemoryRepository.deleteMessages(conversationId);
117+
118+
assertThat(chatMemoryRepository.findMessages(conversationId)).isEmpty();
119+
120+
var multipleMessages = List.<Message>of(new UserMessage("Message from the user 1"),
121+
new AssistantMessage("Message from the assistant 1"));
122+
123+
chatMemoryRepository.saveMessages(conversationId, multipleMessages);
124+
125+
assertThat(chatMemoryRepository.findMessages(conversationId)).hasSize(multipleMessages.size());
126+
assertThat(chatMemoryRepository.findMessages(conversationId)).isEqualTo(multipleMessages);
127+
});
128+
}
129+
130+
@Test
131+
void useAutoConfiguredChatMemoryWithJdbc() {
132+
this.contextRunner.withConfiguration(AutoConfigurations.of(ChatMemoryAutoConfiguration.class)).run(context -> {
133+
assertThat(context).hasSingleBean(ChatMemory.class);
134+
assertThat(context).hasSingleBean(JdbcChatMemoryRepository.class);
135+
136+
var chatMemory = context.getBean(ChatMemory.class);
137+
var conversationId = UUID.randomUUID().toString();
138+
var userMessage = new UserMessage("Message from the user");
139+
140+
chatMemory.add(conversationId, userMessage);
141+
142+
assertThat(chatMemory.get(conversationId)).hasSize(1);
143+
assertThat(chatMemory.get(conversationId)).isEqualTo(List.of(userMessage));
144+
145+
chatMemory.clear(conversationId);
146+
147+
assertThat(chatMemory.get(conversationId)).isEmpty();
148+
149+
var multipleMessages = List.<Message>of(new UserMessage("Message from the user 1"),
150+
new AssistantMessage("Message from the assistant 1"));
151+
152+
chatMemory.add(conversationId, multipleMessages);
153+
154+
assertThat(chatMemory.get(conversationId)).hasSize(multipleMessages.size());
155+
assertThat(chatMemory.get(conversationId)).isEqualTo(multipleMessages);
156+
});
157+
}
158+
99159
}

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-neo4j/pom.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,12 @@
2929
<version>${project.parent.version}</version>
3030
</dependency>
3131

32+
<dependency>
33+
<groupId>org.springframework.ai</groupId>
34+
<artifactId>spring-ai-autoconfigure-model-chat-memory</artifactId>
35+
<version>${project.parent.version}</version>
36+
</dependency>
37+
3238
<!-- Boot dependencies -->
3339
<dependency>
3440
<groupId>org.springframework.boot</groupId>

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-neo4j/src/main/java/org/springframework/ai/model/chat/memory/neo4j/autoconfigure/Neo4jChatMemoryAutoConfiguration.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemory;
2222
import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryConfig;
23+
import org.springframework.ai.model.chat.memory.autoconfigure.ChatMemoryAutoConfiguration;
2324
import org.springframework.boot.autoconfigure.AutoConfiguration;
2425
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
2526
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
@@ -33,7 +34,7 @@
3334
* @author Enrico Rampazzo
3435
* @since 1.0.0
3536
*/
36-
@AutoConfiguration(after = Neo4jAutoConfiguration.class)
37+
@AutoConfiguration(after = Neo4jAutoConfiguration.class, before = ChatMemoryAutoConfiguration.class)
3738
@ConditionalOnClass({ Neo4jChatMemory.class, Driver.class })
3839
@EnableConfigurationProperties(Neo4jChatMemoryProperties.class)
3940
public class Neo4jChatMemoryAutoConfiguration {

0 commit comments

Comments
 (0)