Skip to content

Commit 11e3c8f

Browse files
markpollackmarkpollack
committed
refactor(cassandra): Remove CassandraChatMemory implementation
- Removed the CassandraChatMemory class and its usages, now using CassandraChatMemoryRepository. - Updated configuration, repository, and test classes to rely on CassandraChatMemoryRepository. - Cleaned up related configuration and integration tests. - Minor doc update to reflect the removal. Fixes #3090 Signed-off-by: Mark Pollack <[email protected]>
1 parent 08814be commit 11e3c8f

File tree

16 files changed

+64
-166
lines changed

16 files changed

+64
-166
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818

1919
import com.datastax.oss.driver.api.core.CqlSession;
2020

21-
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryConfig;
21+
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryRepositoryConfig;
2222
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryRepository;
2323
import org.springframework.ai.model.chat.memory.autoconfigure.ChatMemoryAutoConfiguration;
2424
import org.springframework.boot.autoconfigure.AutoConfiguration;
@@ -45,7 +45,7 @@ public class CassandraChatMemoryRepositoryAutoConfiguration {
4545
public CassandraChatMemoryRepository cassandraChatMemoryRepository(
4646
CassandraChatMemoryRepositoryProperties properties, CqlSession cqlSession) {
4747

48-
var builder = CassandraChatMemoryConfig.builder().withCqlSession(cqlSession);
48+
var builder = CassandraChatMemoryRepositoryConfig.builder().withCqlSession(cqlSession);
4949

5050
builder = builder.withKeyspaceName(properties.getKeyspace())
5151
.withTableName(properties.getTable())

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
import org.slf4j.Logger;
2222
import org.slf4j.LoggerFactory;
2323

24-
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryConfig;
24+
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryRepositoryConfig;
2525
import org.springframework.boot.context.properties.ConfigurationProperties;
2626
import org.springframework.lang.Nullable;
2727

@@ -39,13 +39,13 @@ public class CassandraChatMemoryRepositoryProperties {
3939

4040
private static final Logger logger = LoggerFactory.getLogger(CassandraChatMemoryRepositoryProperties.class);
4141

42-
private String keyspace = CassandraChatMemoryConfig.DEFAULT_KEYSPACE_NAME;
42+
private String keyspace = CassandraChatMemoryRepositoryConfig.DEFAULT_KEYSPACE_NAME;
4343

44-
private String table = CassandraChatMemoryConfig.DEFAULT_TABLE_NAME;
44+
private String table = CassandraChatMemoryRepositoryConfig.DEFAULT_TABLE_NAME;
4545

46-
private String assistantColumn = CassandraChatMemoryConfig.DEFAULT_ASSISTANT_COLUMN_NAME;
46+
private String assistantColumn = CassandraChatMemoryRepositoryConfig.DEFAULT_ASSISTANT_COLUMN_NAME;
4747

48-
private String userColumn = CassandraChatMemoryConfig.DEFAULT_USER_COLUMN_NAME;
48+
private String userColumn = CassandraChatMemoryRepositoryConfig.DEFAULT_USER_COLUMN_NAME;
4949

5050
private boolean initializeSchema = true;
5151

‎auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cassandra/src/test/java/org/springframework/ai/model/chat/memory/repository/cassandra/autoconfigure/CassandraChatMemoryRepositoryPropertiesTest.java

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020

2121
import org.junit.jupiter.api.Test;
2222

23-
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryConfig;
23+
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryRepositoryConfig;
2424

2525
import static org.assertj.core.api.Assertions.assertThat;
2626

@@ -34,10 +34,11 @@ class CassandraChatMemoryRepositoryPropertiesTest {
3434
@Test
3535
void defaultValues() {
3636
var props = new CassandraChatMemoryRepositoryProperties();
37-
assertThat(props.getKeyspace()).isEqualTo(CassandraChatMemoryConfig.DEFAULT_KEYSPACE_NAME);
38-
assertThat(props.getTable()).isEqualTo(CassandraChatMemoryConfig.DEFAULT_TABLE_NAME);
39-
assertThat(props.getAssistantColumn()).isEqualTo(CassandraChatMemoryConfig.DEFAULT_ASSISTANT_COLUMN_NAME);
40-
assertThat(props.getUserColumn()).isEqualTo(CassandraChatMemoryConfig.DEFAULT_USER_COLUMN_NAME);
37+
assertThat(props.getKeyspace()).isEqualTo(CassandraChatMemoryRepositoryConfig.DEFAULT_KEYSPACE_NAME);
38+
assertThat(props.getTable()).isEqualTo(CassandraChatMemoryRepositoryConfig.DEFAULT_TABLE_NAME);
39+
assertThat(props.getAssistantColumn())
40+
.isEqualTo(CassandraChatMemoryRepositoryConfig.DEFAULT_ASSISTANT_COLUMN_NAME);
41+
assertThat(props.getUserColumn()).isEqualTo(CassandraChatMemoryRepositoryConfig.DEFAULT_USER_COLUMN_NAME);
4142
assertThat(props.getTimeToLive()).isNull();
4243
assertThat(props.isInitializeSchema()).isTrue();
4344
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818

1919
import org.neo4j.driver.Driver;
2020

21-
import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryConfig;
21+
import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryRepositoryConfig;
2222
import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryRepository;
2323
import org.springframework.ai.model.chat.memory.autoconfigure.ChatMemoryAutoConfiguration;
2424
import org.springframework.boot.autoconfigure.AutoConfiguration;
@@ -44,7 +44,7 @@ public class Neo4jChatMemoryRepositoryAutoConfiguration {
4444
public Neo4jChatMemoryRepository neo4jChatMemoryRepository(Neo4jChatMemoryRepositoryProperties properties,
4545
Driver driver) {
4646

47-
var builder = Neo4jChatMemoryConfig.builder()
47+
var builder = Neo4jChatMemoryRepositoryConfig.builder()
4848
.withMediaLabel(properties.getMediaLabel())
4949
.withMessageLabel(properties.getMessageLabel())
5050
.withMetadataLabel(properties.getMetadataLabel())

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616

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

19-
import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryConfig;
19+
import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryRepositoryConfig;
2020
import org.springframework.boot.context.properties.ConfigurationProperties;
2121

2222
/**
@@ -29,17 +29,17 @@ public class Neo4jChatMemoryRepositoryProperties {
2929

3030
public static final String CONFIG_PREFIX = "spring.ai.chat.memory.neo4j";
3131

32-
private String sessionLabel = Neo4jChatMemoryConfig.DEFAULT_SESSION_LABEL;
32+
private String sessionLabel = Neo4jChatMemoryRepositoryConfig.DEFAULT_SESSION_LABEL;
3333

34-
private String toolCallLabel = Neo4jChatMemoryConfig.DEFAULT_TOOL_CALL_LABEL;
34+
private String toolCallLabel = Neo4jChatMemoryRepositoryConfig.DEFAULT_TOOL_CALL_LABEL;
3535

36-
private String metadataLabel = Neo4jChatMemoryConfig.DEFAULT_METADATA_LABEL;
36+
private String metadataLabel = Neo4jChatMemoryRepositoryConfig.DEFAULT_METADATA_LABEL;
3737

38-
private String messageLabel = Neo4jChatMemoryConfig.DEFAULT_MESSAGE_LABEL;
38+
private String messageLabel = Neo4jChatMemoryRepositoryConfig.DEFAULT_MESSAGE_LABEL;
3939

40-
private String toolResponseLabel = Neo4jChatMemoryConfig.DEFAULT_TOOL_RESPONSE_LABEL;
40+
private String toolResponseLabel = Neo4jChatMemoryRepositoryConfig.DEFAULT_TOOL_RESPONSE_LABEL;
4141

42-
private String mediaLabel = Neo4jChatMemoryConfig.DEFAULT_MEDIA_LABEL;
42+
private String mediaLabel = Neo4jChatMemoryRepositoryConfig.DEFAULT_MEDIA_LABEL;
4343

4444
public String getSessionLabel() {
4545
return this.sessionLabel;

‎auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/model/chat/memory/repository/neo4j/autoconfigure/Neo4JChatMemoryRepositoryPropertiesTest.java

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818

1919
import org.junit.jupiter.api.Test;
2020

21-
import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryConfig;
21+
import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryRepositoryConfig;
2222

2323
import static org.assertj.core.api.Assertions.assertThat;
2424

@@ -31,12 +31,12 @@ class Neo4JChatMemoryRepositoryPropertiesTest {
3131
@Test
3232
void defaultValues() {
3333
var props = new Neo4jChatMemoryRepositoryProperties();
34-
assertThat(props.getMediaLabel()).isEqualTo(Neo4jChatMemoryConfig.DEFAULT_MEDIA_LABEL);
35-
assertThat(props.getMessageLabel()).isEqualTo(Neo4jChatMemoryConfig.DEFAULT_MESSAGE_LABEL);
36-
assertThat(props.getMetadataLabel()).isEqualTo(Neo4jChatMemoryConfig.DEFAULT_METADATA_LABEL);
37-
assertThat(props.getSessionLabel()).isEqualTo(Neo4jChatMemoryConfig.DEFAULT_SESSION_LABEL);
38-
assertThat(props.getToolCallLabel()).isEqualTo(Neo4jChatMemoryConfig.DEFAULT_TOOL_CALL_LABEL);
39-
assertThat(props.getToolResponseLabel()).isEqualTo(Neo4jChatMemoryConfig.DEFAULT_TOOL_RESPONSE_LABEL);
34+
assertThat(props.getMediaLabel()).isEqualTo(Neo4jChatMemoryRepositoryConfig.DEFAULT_MEDIA_LABEL);
35+
assertThat(props.getMessageLabel()).isEqualTo(Neo4jChatMemoryRepositoryConfig.DEFAULT_MESSAGE_LABEL);
36+
assertThat(props.getMetadataLabel()).isEqualTo(Neo4jChatMemoryRepositoryConfig.DEFAULT_METADATA_LABEL);
37+
assertThat(props.getSessionLabel()).isEqualTo(Neo4jChatMemoryRepositoryConfig.DEFAULT_SESSION_LABEL);
38+
assertThat(props.getToolCallLabel()).isEqualTo(Neo4jChatMemoryRepositoryConfig.DEFAULT_TOOL_CALL_LABEL);
39+
assertThat(props.getToolResponseLabel()).isEqualTo(Neo4jChatMemoryRepositoryConfig.DEFAULT_TOOL_RESPONSE_LABEL);
4040
}
4141

4242
}

‎auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/model/chat/memory/repository/neo4j/autoconfigure/Neo4jChatMemoryRepositoryAutoConfigurationIT.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
import org.testcontainers.junit.jupiter.Testcontainers;
3232
import org.testcontainers.utility.DockerImageName;
3333

34-
import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryConfig;
34+
import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryRepositoryConfig;
3535
import org.springframework.ai.chat.messages.AssistantMessage;
3636
import org.springframework.ai.chat.messages.Message;
3737
import org.springframework.ai.chat.messages.SystemMessage;
@@ -151,7 +151,7 @@ void setCustomConfiguration() {
151151
propertyBase.formatted("medialabel", mediaLabel))
152152
.run(context -> {
153153
Neo4jChatMemoryRepository chatMemory = context.getBean(Neo4jChatMemoryRepository.class);
154-
Neo4jChatMemoryConfig config = chatMemory.getConfig();
154+
Neo4jChatMemoryRepositoryConfig config = chatMemory.getConfig();
155155
assertThat(config.getMessageLabel()).isEqualTo(messageLabel);
156156
assertThat(config.getMediaLabel()).isEqualTo(mediaLabel);
157157
assertThat(config.getMetadataLabel()).isEqualTo(metadataLabel);

‎memory/spring-ai-model-chat-memory-cassandra/src/main/java/org/springframework/ai/chat/memory/cassandra/CassandraChatMemory.java

Lines changed: 0 additions & 71 deletions
This file was deleted.

‎memory/spring-ai-model-chat-memory-cassandra/src/main/java/org/springframework/ai/chat/memory/cassandra/CassandraChatMemoryRepository.java

Lines changed: 11 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,6 @@
3939
import org.springframework.ai.chat.messages.UserMessage;
4040
import org.springframework.util.Assert;
4141

42-
import static org.springframework.ai.chat.messages.MessageType.ASSISTANT;
43-
import static org.springframework.ai.chat.messages.MessageType.USER;
44-
4542
/**
4643
* An implementation of {@link ChatMemoryRepository} for Apache Cassandra.
4744
*
@@ -53,7 +50,7 @@ public class CassandraChatMemoryRepository implements ChatMemoryRepository {
5350
public static final String CONVERSATION_TS = CassandraChatMemoryRepository.class.getSimpleName()
5451
+ "_message_timestamp";
5552

56-
final CassandraChatMemoryConfig conf;
53+
final CassandraChatMemoryRepositoryConfig conf;
5754

5855
private final PreparedStatement allStmt;
5956

@@ -65,7 +62,7 @@ public class CassandraChatMemoryRepository implements ChatMemoryRepository {
6562

6663
private final PreparedStatement deleteStmt;
6764

68-
private CassandraChatMemoryRepository(CassandraChatMemoryConfig conf) {
65+
private CassandraChatMemoryRepository(CassandraChatMemoryRepositoryConfig conf) {
6966
Assert.notNull(conf, "conf cannot be null");
7067
this.conf = conf;
7168
this.conf.ensureSchemaExists();
@@ -76,7 +73,7 @@ private CassandraChatMemoryRepository(CassandraChatMemoryConfig conf) {
7673
this.deleteStmt = prepareDeleteStmt();
7774
}
7875

79-
public static CassandraChatMemoryRepository create(CassandraChatMemoryConfig conf) {
76+
public static CassandraChatMemoryRepository create(CassandraChatMemoryRepositoryConfig conf) {
8077
return new CassandraChatMemoryRepository(conf);
8178
}
8279

@@ -91,7 +88,7 @@ public List<String> findConversationIds() {
9188
emptyQuery = true;
9289
for (Row r : this.conf.session.execute(stmt)) {
9390
emptyQuery = false;
94-
conversationIds.add(r.getString(CassandraChatMemoryConfig.DEFAULT_SESSION_ID_NAME));
91+
conversationIds.add(r.getString(CassandraChatMemoryRepositoryConfig.DEFAULT_SESSION_ID_NAME));
9592
token = r.getLong("t");
9693
}
9794
}
@@ -106,7 +103,7 @@ public List<Message> findByConversationId(String conversationId) {
106103
BoundStatementBuilder builder = this.getStmt.boundStatementBuilder();
107104

108105
for (int k = 0; k < primaryKeys.size(); ++k) {
109-
CassandraChatMemoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
106+
CassandraChatMemoryRepositoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
110107
builder = builder.set(keyColumn.name(), primaryKeys.get(k), keyColumn.javaType());
111108
}
112109

@@ -155,13 +152,13 @@ void save(String conversationId, Message msg) {
155152
BoundStatementBuilder builder = stmt.boundStatementBuilder();
156153

157154
for (int k = 0; k < primaryKeys.size(); ++k) {
158-
CassandraChatMemoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
155+
CassandraChatMemoryRepositoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
159156
builder = builder.set(keyColumn.name(), primaryKeys.get(k), keyColumn.javaType());
160157
}
161158

162159
Instant instant = (Instant) msg.getMetadata().get(CONVERSATION_TS);
163160

164-
builder = builder.setInstant(CassandraChatMemoryConfig.DEFAULT_EXCHANGE_ID_NAME, instant)
161+
builder = builder.setInstant(CassandraChatMemoryRepositoryConfig.DEFAULT_EXCHANGE_ID_NAME, instant)
165162
.setString("message", msg.getText());
166163

167164
this.conf.session.execute(builder.build());
@@ -175,7 +172,7 @@ public void deleteByConversationId(String conversationId) {
175172
BoundStatementBuilder builder = this.deleteStmt.boundStatementBuilder();
176173

177174
for (int k = 0; k < primaryKeys.size(); ++k) {
178-
CassandraChatMemoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
175+
CassandraChatMemoryRepositoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
179176
builder = builder.set(keyColumn.name(), primaryKeys.get(k), keyColumn.javaType());
180177
}
181178

@@ -198,10 +195,10 @@ private PreparedStatement prepareAddStmt(String column) {
198195
private PreparedStatement prepareAllStatement() {
199196
Select stmt = QueryBuilder.selectFrom(this.conf.schema.keyspace(), this.conf.schema.table())
200197
.distinct()
201-
.raw(String.format("token(%s)", CassandraChatMemoryConfig.DEFAULT_SESSION_ID_NAME))
198+
.raw(String.format("token(%s)", CassandraChatMemoryRepositoryConfig.DEFAULT_SESSION_ID_NAME))
202199
.as("t")
203-
.column(CassandraChatMemoryConfig.DEFAULT_SESSION_ID_NAME)
204-
.whereToken(CassandraChatMemoryConfig.DEFAULT_SESSION_ID_NAME)
200+
.column(CassandraChatMemoryRepositoryConfig.DEFAULT_SESSION_ID_NAME)
201+
.whereToken(CassandraChatMemoryRepositoryConfig.DEFAULT_SESSION_ID_NAME)
205202
.isGreaterThan(QueryBuilder.bindMarker("after_token"))
206203
.limit(10000);
207204

‎memory/spring-ai-model-chat-memory-cassandra/src/main/java/org/springframework/ai/chat/memory/cassandra/CassandraChatMemoryConfig.java renamed to ‎memory/spring-ai-model-chat-memory-cassandra/src/main/java/org/springframework/ai/chat/memory/cassandra/CassandraChatMemoryRepositoryConfig.java

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@
4949
* @author Mick Semb Wever
5050
* @since 1.0.0
5151
*/
52-
public final class CassandraChatMemoryConfig {
52+
public final class CassandraChatMemoryRepositoryConfig {
5353

5454
public static final String DEFAULT_KEYSPACE_NAME = "springframework";
5555

@@ -65,7 +65,7 @@ public final class CassandraChatMemoryConfig {
6565

6666
public static final String DEFAULT_USER_COLUMN_NAME = "user";
6767

68-
private static final Logger logger = LoggerFactory.getLogger(CassandraChatMemoryConfig.class);
68+
private static final Logger logger = LoggerFactory.getLogger(CassandraChatMemoryRepositoryConfig.class);
6969

7070
final CqlSession session;
7171

@@ -81,7 +81,7 @@ public final class CassandraChatMemoryConfig {
8181

8282
private final boolean disallowSchemaChanges;
8383

84-
private CassandraChatMemoryConfig(Builder builder) {
84+
private CassandraChatMemoryRepositoryConfig(Builder builder) {
8585
this.session = builder.session;
8686
this.schema = new Schema(builder.keyspace, builder.table, builder.partitionKeys, builder.clusteringKeys);
8787
this.assistantColumn = builder.assistantColumn;
@@ -316,7 +316,7 @@ public Builder withChatExchangeToPrimaryKeyTranslator(SessionIdToPrimaryKeysTran
316316
return this;
317317
}
318318

319-
public CassandraChatMemoryConfig build() {
319+
public CassandraChatMemoryRepositoryConfig build() {
320320

321321
int primaryKeyColumns = this.partitionKeys.size() + this.clusteringKeys.size();
322322
int primaryKeysToBind = this.primaryKeyTranslator.apply(UUID.randomUUID().toString()).size();
@@ -328,7 +328,7 @@ public CassandraChatMemoryConfig build() {
328328
this.clusteringKeys.get(this.clusteringKeys.size() - 1).name().equals(DEFAULT_EXCHANGE_ID_NAME),
329329
"last clustering key must be the exchangeIdColumn");
330330

331-
return new CassandraChatMemoryConfig(this);
331+
return new CassandraChatMemoryRepositoryConfig(this);
332332
}
333333

334334
}

0 commit comments

Comments
 (0)