refactor(cassandra): Remove CassandraChatMemory implementation · spring-projects/spring-ai@11e3c8f (original) (raw)

`@@ -39,9 +39,6 @@

`

39

39

`import org.springframework.ai.chat.messages.UserMessage;

`

40

40

`import org.springframework.util.Assert;

`

41

41

``

42

``

`-

import static org.springframework.ai.chat.messages.MessageType.ASSISTANT;

`

43

``

`-

import static org.springframework.ai.chat.messages.MessageType.USER;

`

44

``

-

45

42

`/**

`

46

43

` * An implementation of {@link ChatMemoryRepository} for Apache Cassandra.

`

47

44

` *

`

`@@ -53,7 +50,7 @@ public class CassandraChatMemoryRepository implements ChatMemoryRepository {

`

53

50

`public static final String CONVERSATION_TS = CassandraChatMemoryRepository.class.getSimpleName()

`

54

51

` + "_message_timestamp";

`

55

52

``

56

``

`-

final CassandraChatMemoryConfig conf;

`

``

53

`+

final CassandraChatMemoryRepositoryConfig conf;

`

57

54

``

58

55

`private final PreparedStatement allStmt;

`

59

56

``

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

`

65

62

``

66

63

`private final PreparedStatement deleteStmt;

`

67

64

``

68

``

`-

private CassandraChatMemoryRepository(CassandraChatMemoryConfig conf) {

`

``

65

`+

private CassandraChatMemoryRepository(CassandraChatMemoryRepositoryConfig conf) {

`

69

66

`Assert.notNull(conf, "conf cannot be null");

`

70

67

`this.conf = conf;

`

71

68

`this.conf.ensureSchemaExists();

`

`@@ -76,7 +73,7 @@ private CassandraChatMemoryRepository(CassandraChatMemoryConfig conf) {

`

76

73

`this.deleteStmt = prepareDeleteStmt();

`

77

74

` }

`

78

75

``

79

``

`-

public static CassandraChatMemoryRepository create(CassandraChatMemoryConfig conf) {

`

``

76

`+

public static CassandraChatMemoryRepository create(CassandraChatMemoryRepositoryConfig conf) {

`

80

77

`return new CassandraChatMemoryRepository(conf);

`

81

78

` }

`

82

79

``

`@@ -91,7 +88,7 @@ public List findConversationIds() {

`

91

88

`emptyQuery = true;

`

92

89

`for (Row r : this.conf.session.execute(stmt)) {

`

93

90

`emptyQuery = false;

`

94

``

`-

conversationIds.add(r.getString(CassandraChatMemoryConfig.DEFAULT_SESSION_ID_NAME));

`

``

91

`+

conversationIds.add(r.getString(CassandraChatMemoryRepositoryConfig.DEFAULT_SESSION_ID_NAME));

`

95

92

`token = r.getLong("t");

`

96

93

` }

`

97

94

` }

`

`@@ -106,7 +103,7 @@ public List findByConversationId(String conversationId) {

`

106

103

`BoundStatementBuilder builder = this.getStmt.boundStatementBuilder();

`

107

104

``

108

105

`for (int k = 0; k < primaryKeys.size(); ++k) {

`

109

``

`-

CassandraChatMemoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);

`

``

106

`+

CassandraChatMemoryRepositoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);

`

110

107

`builder = builder.set(keyColumn.name(), primaryKeys.get(k), keyColumn.javaType());

`

111

108

` }

`

112

109

``

`@@ -155,13 +152,13 @@ void save(String conversationId, Message msg) {

`

155

152

`BoundStatementBuilder builder = stmt.boundStatementBuilder();

`

156

153

``

157

154

`for (int k = 0; k < primaryKeys.size(); ++k) {

`

158

``

`-

CassandraChatMemoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);

`

``

155

`+

CassandraChatMemoryRepositoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);

`

159

156

`builder = builder.set(keyColumn.name(), primaryKeys.get(k), keyColumn.javaType());

`

160

157

` }

`

161

158

``

162

159

`Instant instant = (Instant) msg.getMetadata().get(CONVERSATION_TS);

`

163

160

``

164

``

`-

builder = builder.setInstant(CassandraChatMemoryConfig.DEFAULT_EXCHANGE_ID_NAME, instant)

`

``

161

`+

builder = builder.setInstant(CassandraChatMemoryRepositoryConfig.DEFAULT_EXCHANGE_ID_NAME, instant)

`

165

162

` .setString("message", msg.getText());

`

166

163

``

167

164

`this.conf.session.execute(builder.build());

`

`@@ -175,7 +172,7 @@ public void deleteByConversationId(String conversationId) {

`

175

172

`BoundStatementBuilder builder = this.deleteStmt.boundStatementBuilder();

`

176

173

``

177

174

`for (int k = 0; k < primaryKeys.size(); ++k) {

`

178

``

`-

CassandraChatMemoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);

`

``

175

`+

CassandraChatMemoryRepositoryConfig.SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);

`

179

176

`builder = builder.set(keyColumn.name(), primaryKeys.get(k), keyColumn.javaType());

`

180

177

` }

`

181

178

``

`@@ -198,10 +195,10 @@ private PreparedStatement prepareAddStmt(String column) {

`

198

195

`private PreparedStatement prepareAllStatement() {

`

199

196

`Select stmt = QueryBuilder.selectFrom(this.conf.schema.keyspace(), this.conf.schema.table())

`

200

197

` .distinct()

`

201

``

`-

.raw(String.format("token(%s)", CassandraChatMemoryConfig.DEFAULT_SESSION_ID_NAME))

`

``

198

`+

.raw(String.format("token(%s)", CassandraChatMemoryRepositoryConfig.DEFAULT_SESSION_ID_NAME))

`

202

199

` .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)

`

205

202

` .isGreaterThan(QueryBuilder.bindMarker("after_token"))

`

206

203

` .limit(10000);

`

207

204

``