Implement CassandraChatMemory

- provides a flexible schema, can be combined with a vector store, and supports time-to-live rows.
 - fix initialize-schema docs and so that it actually works.
 - move CommonVectorStoreProperties to .vectorstore. package
 - add CassandraAutoConfiguration to the AutoConfiguration.imports
This commit is contained in:
mck
2024-06-11 17:13:12 +02:00
committed by Christian Tzolov
parent cf785e03c2
commit 4aacab020b
35 changed files with 1089 additions and 68 deletions

4
.gitignore vendored
View File

@@ -38,4 +38,6 @@ package.json
shell.log
.profiler
.profiler
/spring-ai-spring-boot-autoconfigure/nbproject/
/vector-stores/spring-ai-cassandra-store/nbproject/

View File

@@ -170,7 +170,7 @@
<protobuf-java.version>3.25.2</protobuf-java.version>
<!-- readers/writer/stores dependencies-->
<cassandra.java-driver.version>4.18.0</cassandra.java-driver.version>
<cassandra.java-driver.version>4.18.1</cassandra.java-driver.version>
<pdfbox.version>3.0.1</pdfbox.version>
<pgvector.version>0.1.4</pgvector.version>
<sap.hanadb.version>2.20.11</sap.hanadb.version>

View File

@@ -26,7 +26,7 @@ import org.springframework.ai.chat.messages.Message;
* conversation, and clear the conversation history.
*
* @author Christian Tzolov
* @since 1.0.0 M1
* @since 1.0.0
*/
public interface ChatMemory {

View File

@@ -369,7 +369,13 @@ The `FILTER_EXPRESSION` parameter allows you to dynamically filter the search re
The interface `ChatMemory` represents a storage for chat conversation history. It provides methods to add messages to a
* conversation, retrieve messages from a conversation, and clear the conversation history.
There is one implementation `InMemoryChatMemory` that provides in-memory storage for chat conversation history.
There are two implementations `InMemoryChatMemory` and `CassandraChatMemory` that provides storage for chat conversation history, in-memory and persisted with time-to-live correspondingly.
To create a CassandraChatMemory with time-to-live
[source,java]
----
CassandraChatMemory.create(CassandraChatMemoryConfig.builder().withTimeToLive(Duration.ofDays(1)).build());
----
Two advisor implementations use the `ChatMemory` interface to advice the prompt with conversation history which differ in the details of how the memory is added to the prompt

View File

@@ -20,14 +20,11 @@ This Spring AI Vector Store is designed to work for both brand-new RAG applicati
The store can also be used for non-RAG use-cases in an existing database, e.g. semantic searches, geo-proximity searches, etc.
The vector store implementation can initialize the requisite schema for you, but you must opt-in by specifying the `initializeSchema` boolean in the appropriate constructor or by setting `...initialize-schema=true` in the `application.properties` file.
NOTE: this is a breaking change! In earlier versions of Spring AI, this schema initialization happened by default.
The store will automatically create, or enhance, the schema as needed according to its configuration. If you don't want the schema modifications, configure the store with `disallowSchemaChanges`.
When using spring-boot-autoconfigure `disallowSchemaChanges` defaults to true, per Spring Boot standards, and you must opt-in to schema creation/modifications by specifying the `initializeSchema` boolean in the appropriate constructor or by setting `...initialize-schema=true` in the `application.properties` file.
== What is JVector ?
link:https://github.com/jbellis/jvector[JVector] is a pure Java embedded vector search engine.

View File

@@ -0,0 +1,34 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.chat.memory;
/**
* @author Mick Semb Wever
* @since 1.0.0
*/
public class CommonChatMemoryProperties {
private boolean initializeSchema = true;
public boolean isInitializeSchema() {
return initializeSchema;
}
public void setInitializeSchema(boolean initializeSchema) {
this.initializeSchema = initializeSchema;
}
}

View File

@@ -0,0 +1,59 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.chat.memory.cassandra;
import com.datastax.oss.driver.api.core.CqlSession;
import org.springframework.ai.chat.memory.CassandraChatMemory;
import org.springframework.ai.chat.memory.CassandraChatMemoryConfig;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.cassandra.CassandraAutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
/**
* @author Mick Semb Wever
* @since 1.0.0
*/
@AutoConfiguration(after = CassandraAutoConfiguration.class)
@ConditionalOnClass({ CassandraChatMemory.class, CqlSession.class })
@EnableConfigurationProperties(CassandraChatMemoryProperties.class)
public class CassandraChatMemoryAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public CassandraChatMemory chatMemory(CassandraChatMemoryProperties properties, CqlSession cqlSession) {
var builder = CassandraChatMemoryConfig.builder().withCqlSession(cqlSession);
builder = builder.withKeyspaceName(properties.getKeyspace())
.withTableName(properties.getTable())
.withAssistantColumnName(properties.getAssistantColumn())
.withUserColumnName(properties.getUserColumn());
if (!properties.isInitializeSchema()) {
builder = builder.disallowSchemaChanges();
}
if (null != properties.getTimeToLiveSeconds()) {
builder = builder.withTimeToLive(properties.getTimeToLiveSeconds());
}
return CassandraChatMemory.create(builder.build());
}
}

View File

@@ -0,0 +1,89 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.chat.memory.cassandra;
import java.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.autoconfigure.chat.memory.CommonChatMemoryProperties;
import org.springframework.ai.chat.memory.CassandraChatMemoryConfig;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.lang.Nullable;
/**
* @author Mick Semb Wever
* @since 1.0.0
*/
@ConfigurationProperties(CassandraChatMemoryProperties.CONFIG_PREFIX)
public class CassandraChatMemoryProperties extends CommonChatMemoryProperties {
public static final String CONFIG_PREFIX = "spring.ai.chat.memory.cassandra";
private static final Logger logger = LoggerFactory.getLogger(CassandraChatMemoryProperties.class);
private String keyspace = CassandraChatMemoryConfig.DEFAULT_KEYSPACE_NAME;
private String table = CassandraChatMemoryConfig.DEFAULT_TABLE_NAME;
private String assistantColumn = CassandraChatMemoryConfig.DEFAULT_ASSISTANT_COLUMN_NAME;
private String userColumn = CassandraChatMemoryConfig.DEFAULT_USER_COLUMN_NAME;
private Duration timeToLiveSeconds = null;
public String getKeyspace() {
return this.keyspace;
}
public void setKeyspace(String keyspace) {
this.keyspace = keyspace;
}
public String getTable() {
return this.table;
}
public void setTable(String table) {
this.table = table;
}
public String getAssistantColumn() {
return assistantColumn;
}
public void setAssistantColumn(String assistantColumn) {
this.assistantColumn = assistantColumn;
}
public String getUserColumn() {
return userColumn;
}
public void setUserColumn(String userColumn) {
this.userColumn = userColumn;
}
@Nullable
public Duration getTimeToLiveSeconds() {
return timeToLiveSeconds;
}
public void setTimeToLiveSeconds(Duration timeToLiveSeconds) {
this.timeToLiveSeconds = timeToLiveSeconds;
}
}

View File

@@ -15,7 +15,7 @@
*/
package org.springframework.ai.autoconfigure.vectorstore.azure;
import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.azure.AzureVectorStore;
import org.springframework.boot.context.properties.ConfigurationProperties;

View File

@@ -54,14 +54,14 @@ public class CassandraVectorStoreAutoConfiguration {
.withIndexName(properties.getIndexName())
.withFixedThreadPoolExecutorSize(properties.getFixedThreadPoolExecutorSize());
if (properties.getDisallowSchemaCreation()) {
if (!properties.isInitializeSchema()) {
builder = builder.disallowSchemaChanges();
}
if (properties.getReturnEmbeddings()) {
builder = builder.returnEmbeddings();
}
return new CassandraVectorStore(builder.build(), embeddingModel);
return CassandraVectorStore.create(builder.build(), embeddingModel);
}
@Bean

View File

@@ -16,7 +16,10 @@
package org.springframework.ai.autoconfigure.vectorstore.cassandra;
import com.google.api.client.util.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig;
import org.springframework.boot.context.properties.ConfigurationProperties;
@@ -25,10 +28,12 @@ import org.springframework.boot.context.properties.ConfigurationProperties;
* @since 1.0.0
*/
@ConfigurationProperties(CassandraVectorStoreProperties.CONFIG_PREFIX)
public class CassandraVectorStoreProperties {
public class CassandraVectorStoreProperties extends CommonVectorStoreProperties {
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.cassandra";
private static final Logger logger = LoggerFactory.getLogger(CassandraVectorStoreProperties.class);
private String keyspace = CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME;
private String table = CassandraVectorStoreConfig.DEFAULT_TABLE_NAME;
@@ -39,8 +44,6 @@ public class CassandraVectorStoreProperties {
private String embeddingColumnName = CassandraVectorStoreConfig.DEFAULT_EMBEDDING_COLUMN_NAME;
private boolean disallowSchemaChanges = false;
private boolean returnEmbeddings = false;
private int fixedThreadPoolExecutorSize = CassandraVectorStoreConfig.DEFAULT_ADD_CONCURRENCY;
@@ -85,12 +88,16 @@ public class CassandraVectorStoreProperties {
this.embeddingColumnName = embeddingColumnName;
}
@Deprecated
public boolean getDisallowSchemaCreation() {
return this.disallowSchemaChanges;
logger.warn("getDisallowSchemaCreation() is deprecated, use isInitializeSchema()");
return !super.isInitializeSchema();
}
@Deprecated
public void setDisallowSchemaCreation(boolean disallowSchemaCreation) {
this.disallowSchemaChanges = disallowSchemaCreation;
logger.warn("setDisallowSchemaCreation(boolean) is deprecated, use setInitializeSchema(boolean)");
super.setInitializeSchema(!disallowSchemaCreation);
}
public boolean getReturnEmbeddings() {

View File

@@ -15,7 +15,7 @@
*/
package org.springframework.ai.autoconfigure.vectorstore.chroma;
import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.ChromaVectorStore;
import org.springframework.boot.context.properties.ConfigurationProperties;

View File

@@ -15,7 +15,7 @@
*/
package org.springframework.ai.autoconfigure.vectorstore.elasticsearch;
import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.SimilarityFunction;
import org.springframework.boot.context.properties.ConfigurationProperties;

View File

@@ -15,7 +15,7 @@
*/
package org.springframework.ai.autoconfigure.vectorstore.milvus;
import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.MilvusVectorStore;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.util.Assert;

View File

@@ -15,7 +15,7 @@
*/
package org.springframework.ai.autoconfigure.vectorstore.mongo;
import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**

View File

@@ -15,7 +15,7 @@
*/
package org.springframework.ai.autoconfigure.vectorstore.neo4j;
import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.Neo4jVectorStore;
import org.springframework.boot.context.properties.ConfigurationProperties;

View File

@@ -15,7 +15,7 @@
*/
package org.springframework.ai.autoconfigure.vectorstore.oracle;
import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.OracleVectorStore;
import org.springframework.boot.context.properties.ConfigurationProperties;

View File

@@ -15,7 +15,7 @@
*/
package org.springframework.ai.autoconfigure.vectorstore.pgvector;
import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.PgVectorStore.PgDistanceType;
import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType;

View File

@@ -15,7 +15,7 @@
*/
package org.springframework.ai.autoconfigure.vectorstore.qdrant;
import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.boot.context.properties.ConfigurationProperties;

View File

@@ -15,7 +15,7 @@
*/
package org.springframework.ai.autoconfigure.vectorstore.redis;
import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**

View File

@@ -17,7 +17,7 @@ package org.springframework.ai.autoconfigure.vectorstore.weaviate;
import java.util.Map;
import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.WeaviateVectorStore.WeaviateVectorStoreConfig;
import org.springframework.ai.vectorstore.WeaviateVectorStore.WeaviateVectorStoreConfig.ConsistentLevel;
import org.springframework.ai.vectorstore.WeaviateVectorStore.WeaviateVectorStoreConfig.MetadataField;

View File

@@ -41,4 +41,5 @@ org.springframework.ai.autoconfigure.vectorstore.opensearch.OpenSearchVectorStor
org.springframework.ai.autoconfigure.moonshot.MoonshotAutoConfiguration
org.springframework.ai.autoconfigure.qianfan.QianFanAutoConfiguration
org.springframework.ai.autoconfigure.minimax.MiniMaxAutoConfiguration
org.springframework.ai.autoconfigure.vertexai.embedding.VertexAiEmbeddingAutoConfiguration
org.springframework.ai.autoconfigure.vertexai.embedding.VertexAiEmbeddingAutoConfiguration
org.springframework.ai.autoconfigure.chat.memory.cassandra.CassandraAutoConfiguration

View File

@@ -0,0 +1,97 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.chat.memory.cassandra;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.memory.CassandraChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.autoconfigure.cassandra.CassandraAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.testcontainers.containers.CassandraContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.utility.DockerImageName;
import com.datastax.driver.core.utils.UUIDs;
/**
* @author Mick Semb Wever
* @since 1.0.0
*/
@Testcontainers
class CassandraChatMemoryAutoConfigurationIT {
static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse("cassandra");
@Container
static CassandraContainer cassandraContainer = new CassandraContainer(DEFAULT_IMAGE_NAME.withTag("5.0"));
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(
AutoConfigurations.of(CassandraChatMemoryAutoConfiguration.class, CassandraAutoConfiguration.class))
.withPropertyValues("spring.ai.chat.memory.cassandra.keyspace=test_autoconfigure");
@Test
void addAndGet() {
contextRunner.withPropertyValues("spring.cassandra.contactPoints=" + getContactPointHost())
.withPropertyValues("spring.cassandra.port=" + getContactPointPort())
.withPropertyValues("spring.cassandra.localDatacenter=" + cassandraContainer.getLocalDatacenter())
.run(context -> {
CassandraChatMemory memory = context.getBean(CassandraChatMemory.class);
String sessionId = UUIDs.timeBased().toString();
assertThat(memory.get(sessionId, Integer.MAX_VALUE)).isEmpty();
memory.add(sessionId, new UserMessage("test question"));
assertThat(memory.get(sessionId, Integer.MAX_VALUE)).hasSize(1);
assertThat(memory.get(sessionId, Integer.MAX_VALUE).get(0).getMessageType())
.isEqualTo(MessageType.USER);
assertThat(memory.get(sessionId, Integer.MAX_VALUE).get(0).getContent()).isEqualTo("test question");
memory.clear(sessionId);
assertThat(memory.get(sessionId, Integer.MAX_VALUE)).isEmpty();
memory.add(sessionId, List.of(new UserMessage("test question"), new AssistantMessage("test answer")));
assertThat(memory.get(sessionId, Integer.MAX_VALUE)).hasSize(2);
assertThat(memory.get(sessionId, Integer.MAX_VALUE).get(1).getMessageType())
.isEqualTo(MessageType.USER);
assertThat(memory.get(sessionId, Integer.MAX_VALUE).get(1).getContent()).isEqualTo("test question");
assertThat(memory.get(sessionId, Integer.MAX_VALUE).get(0).getMessageType())
.isEqualTo(MessageType.ASSISTANT);
assertThat(memory.get(sessionId, Integer.MAX_VALUE).get(0).getContent()).isEqualTo("test answer");
});
}
private String getContactPointHost() {
return cassandraContainer.getContactPoint().getHostString();
}
private String getContactPointPort() {
return String.valueOf(cassandraContainer.getContactPoint().getPort());
}
}

View File

@@ -0,0 +1,61 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.chat.memory.cassandra;
import java.time.Duration;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.memory.CassandraChatMemoryConfig;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Mick Semb Wever
* @since 1.0.0
*/
class CassandraChatMemoryPropertiesTest {
@Test
void defaultValues() {
var props = new CassandraChatMemoryProperties();
assertThat(props.getKeyspace()).isEqualTo(CassandraChatMemoryConfig.DEFAULT_KEYSPACE_NAME);
assertThat(props.getTable()).isEqualTo(CassandraChatMemoryConfig.DEFAULT_TABLE_NAME);
assertThat(props.getAssistantColumn()).isEqualTo(CassandraChatMemoryConfig.DEFAULT_ASSISTANT_COLUMN_NAME);
assertThat(props.getUserColumn()).isEqualTo(CassandraChatMemoryConfig.DEFAULT_USER_COLUMN_NAME);
assertThat(props.getTimeToLiveSeconds()).isNull();
assertThat(props.isInitializeSchema()).isTrue();
}
@Test
void customValues() {
var props = new CassandraChatMemoryProperties();
props.setKeyspace("my_keyspace");
props.setTable("my_table");
props.setAssistantColumn("my_assistant_column");
props.setUserColumn("my_user_column");
props.setTimeToLiveSeconds(Duration.ofDays(1));
props.setInitializeSchema(false);
assertThat(props.getKeyspace()).isEqualTo("my_keyspace");
assertThat(props.getTable()).isEqualTo("my_table");
assertThat(props.getAssistantColumn()).isEqualTo("my_assistant_column");
assertThat(props.getUserColumn()).isEqualTo("my_user_column");
assertThat(props.getTimeToLiveSeconds()).isEqualTo(Duration.ofDays(1));
assertThat(props.isInitializeSchema()).isFalse();
}
}

View File

@@ -0,0 +1,68 @@
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
* See the NOTICE file distributed with this work for additional information
* regarding copyright ownership.
*/
package org.springframework.ai.cassandra;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.querybuilder.SchemaBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.time.Duration;
/**
* @author Mick Semb Wever
* @since 1.0.0
*/
public final class SchemaUtil {
private static final Logger logger = LoggerFactory.getLogger(SchemaUtil.class);
private SchemaUtil() {
}
public static void checkSchemaAgreement(CqlSession session) throws IllegalStateException {
if (!session.checkSchemaAgreement()) {
logger.warn("Waiting for cluster schema agreement, sleeping 10s…");
try {
Thread.sleep(Duration.ofSeconds(10).toMillis());
}
catch (InterruptedException ex) {
Thread.currentThread().interrupt();
throw new IllegalStateException(ex);
}
if (!session.checkSchemaAgreement()) {
logger.error("no cluster schema agreement still, continuing, let's hope this works…");
}
}
}
public static void ensureKeyspaceExists(CqlSession session, String keyspaceName) {
if (session.getMetadata().getKeyspace(keyspaceName).isEmpty()) {
SimpleStatement keyspaceStmt = SchemaBuilder.createKeyspace(keyspaceName)
.ifNotExists()
.withSimpleStrategy(1)
.build();
logger.debug("Executing {}", keyspaceStmt.getQuery());
session.execute(keyspaceStmt);
}
}
}

View File

@@ -0,0 +1,192 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.memory;
import com.datastax.oss.driver.api.core.cql.BoundStatementBuilder;
import com.datastax.oss.driver.api.core.cql.PreparedStatement;
import com.datastax.oss.driver.api.core.cql.Row;
import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
import com.datastax.oss.driver.api.querybuilder.delete.Delete;
import com.datastax.oss.driver.api.querybuilder.delete.DeleteSelection;
import com.datastax.oss.driver.api.querybuilder.insert.InsertInto;
import com.datastax.oss.driver.api.querybuilder.insert.RegularInsert;
import com.datastax.oss.driver.api.querybuilder.select.Select;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import org.springframework.ai.chat.memory.CassandraChatMemoryConfig.SchemaColumn;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;
/**
* Create a CassandraChatMemory like <code>
CassandraChatMemory.create(CassandraChatMemoryConfig.builder().withTimeToLive(Duration.ofDays(1)).build());
</code>
*
* For example @see org.springframework.ai.chat.memory.CassandraChatMemory
*
* @author Mick Semb Wever
* @since 1.0.0
*/
public final class CassandraChatMemory implements ChatMemory {
public static final String CONVERSATION_TS = CassandraChatMemory.class.getSimpleName() + "_message_timestamp";
final CassandraChatMemoryConfig conf;
private final PreparedStatement addUserStmt, addAssistantStmt, getStmt, deleteStmt;
public static CassandraChatMemory create(CassandraChatMemoryConfig conf) {
return new CassandraChatMemory(conf);
}
public CassandraChatMemory(CassandraChatMemoryConfig config) {
this.conf = config;
this.conf.ensureSchemaExists();
this.addUserStmt = prepareAddStmt(this.conf.userColumn);
this.addAssistantStmt = prepareAddStmt(this.conf.assistantColumn);
this.getStmt = prepareGetStatement();
this.deleteStmt = prepareDeleteStmt();
}
@Override
public void add(String conversationId, List<Message> messages) {
final AtomicLong instantSeq = new AtomicLong(Instant.now().toEpochMilli());
messages.forEach((msg) -> {
if (msg.getMetadata().containsKey(CONVERSATION_TS)) {
msg.getMetadata().put(CONVERSATION_TS, Instant.ofEpochMilli(instantSeq.getAndIncrement()));
}
add(conversationId, msg);
});
}
@Override
public void add(String sessionId, Message msg) {
Preconditions.checkArgument(
!msg.getMetadata().containsKey(CONVERSATION_TS)
|| msg.getMetadata().get(CONVERSATION_TS) instanceof Instant,
"messages only accept metadata '%s' entries of type Instant", CONVERSATION_TS);
msg.getMetadata().putIfAbsent(CONVERSATION_TS, Instant.now());
PreparedStatement stmt;
switch (msg.getMessageType()) {
case USER -> stmt = addUserStmt;
case ASSISTANT -> stmt = addAssistantStmt;
default -> throw new IllegalArgumentException("Cant add type " + msg);
}
List<Object> primaryKeys = this.conf.primaryKeyTranslator.apply(sessionId);
BoundStatementBuilder builder = stmt.boundStatementBuilder();
for (int k = 0; k < primaryKeys.size(); ++k) {
SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
builder = builder.set(keyColumn.name(), primaryKeys.get(k), keyColumn.javaType());
}
Instant instant = (Instant) msg.getMetadata().get(CONVERSATION_TS);
builder = builder.setInstant(CassandraChatMemoryConfig.DEFAULT_EXCHANGE_ID_NAME, instant)
.setString("message", msg.getContent());
this.conf.session.execute(builder.build());
}
@Override
public void clear(String sessionId) {
List<Object> primaryKeys = this.conf.primaryKeyTranslator.apply(sessionId);
BoundStatementBuilder builder = deleteStmt.boundStatementBuilder();
for (int k = 0; k < primaryKeys.size(); ++k) {
SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
builder = builder.set(keyColumn.name(), primaryKeys.get(k), keyColumn.javaType());
}
this.conf.session.execute(builder.build());
}
@Override
public List<Message> get(String sessionId, int lastN) {
List<Object> primaryKeys = this.conf.primaryKeyTranslator.apply(sessionId);
BoundStatementBuilder builder = getStmt.boundStatementBuilder().setInt("lastN", lastN);
for (int k = 0; k < primaryKeys.size(); ++k) {
SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
builder = builder.set(keyColumn.name(), primaryKeys.get(k), keyColumn.javaType());
}
List<Message> messages = new ArrayList<>();
for (Row r : this.conf.session.execute(builder.build())) {
String assistant = r.getString(this.conf.assistantColumn);
String user = r.getString(this.conf.userColumn);
if (null != assistant) {
messages.add(new AssistantMessage(assistant));
}
if (null != user) {
messages.add(new UserMessage(user));
}
}
return messages;
}
private PreparedStatement prepareAddStmt(String column) {
RegularInsert stmt = null;
InsertInto stmtStart = QueryBuilder.insertInto(this.conf.schema.keyspace(), this.conf.schema.table());
for (var c : this.conf.schema.partitionKeys()) {
stmt = (null != stmt ? stmt : stmtStart).value(c.name(), QueryBuilder.bindMarker(c.name()));
}
for (var c : this.conf.schema.clusteringKeys()) {
stmt = stmt.value(c.name(), QueryBuilder.bindMarker(c.name()));
}
stmt = stmt.value(column, QueryBuilder.bindMarker("message"));
return this.conf.session.prepare(stmt.build());
}
private PreparedStatement prepareGetStatement() {
Select stmt = QueryBuilder.selectFrom(this.conf.schema.keyspace(), this.conf.schema.table()).all();
for (var c : this.conf.schema.partitionKeys()) {
stmt = stmt.whereColumn(c.name()).isEqualTo(QueryBuilder.bindMarker(c.name()));
}
for (int i = 0; i + 1 < this.conf.schema.clusteringKeys().size(); ++i) {
String columnName = this.conf.schema.clusteringKeys().get(i).name();
stmt = stmt.whereColumn(columnName).isEqualTo(QueryBuilder.bindMarker(columnName));
}
stmt = stmt.limit(QueryBuilder.bindMarker("lastN"));
return this.conf.session.prepare(stmt.build());
}
private PreparedStatement prepareDeleteStmt() {
Delete stmt = null;
DeleteSelection stmtStart = QueryBuilder.deleteFrom(this.conf.schema.keyspace(), this.conf.schema.table());
for (var c : this.conf.schema.partitionKeys()) {
stmt = (null != stmt ? stmt : stmtStart).whereColumn(c.name()).isEqualTo(QueryBuilder.bindMarker(c.name()));
}
for (int i = 0; i + 1 < this.conf.schema.clusteringKeys().size(); ++i) {
String columnName = this.conf.schema.clusteringKeys().get(i).name();
stmt = stmt.whereColumn(columnName).isEqualTo(QueryBuilder.bindMarker(columnName));
}
return this.conf.session.prepare(stmt.build());
}
}

View File

@@ -0,0 +1,332 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.memory;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder;
import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
import com.datastax.oss.driver.api.core.type.DataType;
import com.datastax.oss.driver.api.core.type.DataTypes;
import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
import com.datastax.oss.driver.api.core.type.reflect.GenericType;
import com.datastax.oss.driver.api.querybuilder.SchemaBuilder;
import com.datastax.oss.driver.api.querybuilder.schema.AlterTableAddColumn;
import com.datastax.oss.driver.api.querybuilder.schema.AlterTableAddColumnEnd;
import com.datastax.oss.driver.api.querybuilder.schema.CreateTable;
import com.datastax.oss.driver.api.querybuilder.schema.CreateTableStart;
import com.datastax.oss.driver.api.querybuilder.schema.CreateTableWithOptions;
import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.cassandra.SchemaUtil;
import java.net.InetSocketAddress;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.function.Function;
/**
* @author Mick Semb Wever
* @since 1.0.0
*/
public final class CassandraChatMemoryConfig {
private static final Logger logger = LoggerFactory.getLogger(CassandraChatMemoryConfig.class);
record Schema(String keyspace, String table, List<SchemaColumn> partitionKeys, List<SchemaColumn> clusteringKeys) {
}
public record SchemaColumn(String name, DataType type) {
public GenericType<Object> javaType() {
return CodecRegistry.DEFAULT.codecFor(type).getJavaType();
}
}
/** Given a string sessionId, return the value for each primary key column. */
public interface SessionIdToPrimaryKeysTranslator extends Function<String, List<Object>> {
}
public static final String DEFAULT_KEYSPACE_NAME = "springframework";
public static final String DEFAULT_TABLE_NAME = "ai_chat_memory";
// todo make configurable
public static final String DEFAULT_SESSION_ID_NAME = "session_id";
// todo make configurable
public static final String DEFAULT_EXCHANGE_ID_NAME = "message_timestamp";
public static final String DEFAULT_ASSISTANT_COLUMN_NAME = "assistant";
public static final String DEFAULT_USER_COLUMN_NAME = "user";
final CqlSession session;
final Schema schema;
final String assistantColumn;
final String userColumn;
private final Integer timeToLiveSeconds;
private final boolean disallowSchemaChanges;
final SessionIdToPrimaryKeysTranslator primaryKeyTranslator;
public static Builder builder() {
return new Builder();
}
private CassandraChatMemoryConfig(Builder builder) {
this.session = builder.session;
this.schema = new Schema(builder.keyspace, builder.table, builder.partitionKeys, builder.clusteringKeys);
this.assistantColumn = builder.assistantColumn;
this.userColumn = builder.userColumn;
this.timeToLiveSeconds = builder.timeToLiveSeconds;
this.disallowSchemaChanges = builder.disallowSchemaChanges;
this.primaryKeyTranslator = builder.primaryKeyTranslator;
}
SchemaColumn getPrimaryKeyColumn(int index) {
return index < this.schema.partitionKeys().size() ? this.schema.partitionKeys().get(index)
: this.schema.clusteringKeys().get(index - this.schema.partitionKeys().size());
}
@VisibleForTesting
void dropKeyspace() {
Preconditions.checkState(this.schema.keyspace.startsWith("test_"), "Only test keyspaces can be dropped");
this.session.execute(SchemaBuilder.dropKeyspace(this.schema.keyspace).ifExists().build());
}
public static class Builder {
private CqlSession session = null;
private CqlSessionBuilder sessionBuilder = null;
private String keyspace = DEFAULT_KEYSPACE_NAME;
private String table = DEFAULT_TABLE_NAME;
private List<SchemaColumn> partitionKeys = List.of(new SchemaColumn(DEFAULT_SESSION_ID_NAME, DataTypes.TEXT));
private List<SchemaColumn> clusteringKeys = List
.of(new SchemaColumn(DEFAULT_EXCHANGE_ID_NAME, DataTypes.TIMESTAMP));
private String assistantColumn = DEFAULT_ASSISTANT_COLUMN_NAME;
private String userColumn = DEFAULT_USER_COLUMN_NAME;
private Integer timeToLiveSeconds = null;
private boolean disallowSchemaChanges = false;
private SessionIdToPrimaryKeysTranslator primaryKeyTranslator = (sessionId) -> List.of(sessionId);
private Builder() {
}
public Builder withCqlSession(CqlSession session) {
Preconditions.checkState(null == this.sessionBuilder,
"Cannot call withContactPoint(..) or withLocalDatacenter(..) and this method");
this.session = session;
return this;
}
public Builder addContactPoint(InetSocketAddress contactPoint) {
Preconditions.checkState(null == this.session, "Cannot call withCqlSession(..) and this method");
if (null == this.sessionBuilder) {
this.sessionBuilder = new CqlSessionBuilder();
}
this.sessionBuilder.addContactPoint(contactPoint);
return this;
}
public Builder withLocalDatacenter(String localDC) {
Preconditions.checkState(null == this.session, "Cannot call withCqlSession(..) and this method");
if (null == this.sessionBuilder) {
this.sessionBuilder = new CqlSessionBuilder();
}
this.sessionBuilder.withLocalDatacenter(localDC);
return this;
}
public Builder withKeyspaceName(String keyspace) {
this.keyspace = keyspace;
return this;
}
public Builder withTableName(String table) {
this.table = table;
return this;
}
public Builder withPartitionKeys(List<SchemaColumn> partitionKeys) {
Preconditions.checkArgument(!partitionKeys.isEmpty());
this.partitionKeys = partitionKeys;
return this;
}
public Builder withClusteringKeys(List<SchemaColumn> clusteringKeys) {
Preconditions.checkArgument(!clusteringKeys.isEmpty());
this.clusteringKeys = clusteringKeys;
return this;
}
public Builder withAssistantColumnName(String name) {
this.assistantColumn = name;
return this;
}
public Builder withUserColumnName(String name) {
this.userColumn = name;
return this;
}
/** How long are messages kept for */
public Builder withTimeToLive(Duration timeToLive) {
Preconditions.checkArgument(0 < timeToLive.getSeconds());
this.timeToLiveSeconds = (int) timeToLive.toSeconds();
return this;
}
public Builder disallowSchemaChanges() {
this.disallowSchemaChanges = true;
return this;
}
public Builder withChatExchangeToPrimaryKeyTranslator(SessionIdToPrimaryKeysTranslator primaryKeyTranslator) {
this.primaryKeyTranslator = primaryKeyTranslator;
return this;
}
public CassandraChatMemoryConfig build() {
int primaryKeyColumns = partitionKeys.size() + clusteringKeys.size();
int primaryKeysToBind = this.primaryKeyTranslator.apply(UUID.randomUUID().toString()).size();
Preconditions.checkArgument(primaryKeyColumns == primaryKeysToBind + 1,
"The primaryKeyTranslator must always return one less element than the number of primary keys in total. The last clustering key remains undefined, expecting to be the timestamp for messages within sessionId. The sessionId can map to any primary key column (though it should map to a partition key column).");
Preconditions.checkArgument(
clusteringKeys.get(clusteringKeys.size() - 1).name().equals(DEFAULT_EXCHANGE_ID_NAME),
"last clustering key must be the exchangeIdColumn");
return new CassandraChatMemoryConfig(this);
}
}
void ensureSchemaExists() {
if (!disallowSchemaChanges) {
SchemaUtil.ensureKeyspaceExists(this.session, this.schema.keyspace);
ensureTableExists();
ensureTableColumnsExist();
SchemaUtil.checkSchemaAgreement(this.session);
}
else {
checkSchemaValid();
}
}
void checkSchemaValid() {
Preconditions.checkState(session.getMetadata().getKeyspace(this.schema.keyspace).isPresent(),
"keyspace %s does not exist", this.schema.keyspace);
Preconditions.checkState(
session.getMetadata().getKeyspace(this.schema.keyspace).get().getTable(this.schema.table).isPresent(),
"table %s does not exist");
TableMetadata tableMetadata = session.getMetadata()
.getKeyspace(this.schema.keyspace)
.get()
.getTable(this.schema.table)
.get();
Preconditions.checkState(tableMetadata.getColumn(this.assistantColumn).isPresent(), "column %s does not exist",
this.assistantColumn);
Preconditions.checkState(tableMetadata.getColumn(this.userColumn).isPresent(), "column %s does not exist",
this.userColumn);
}
private void ensureTableExists() {
if (session.getMetadata().getKeyspace(schema.keyspace).get().getTable(this.schema.table).isEmpty()) {
CreateTable createTable = null;
CreateTableStart createTableStart = SchemaBuilder.createTable(this.schema.keyspace, this.schema.table)
.ifNotExists();
for (SchemaColumn partitionKey : this.schema.partitionKeys) {
createTable = (null != createTable ? createTable : createTableStart).withPartitionKey(partitionKey.name,
partitionKey.type);
}
for (SchemaColumn clusteringKey : this.schema.clusteringKeys) {
createTable = createTable.withClusteringColumn(clusteringKey.name, clusteringKey.type);
}
String lastClusteringColumn = this.schema.clusteringKeys.get(this.schema.clusteringKeys.size() - 1).name();
CreateTableWithOptions createTableWithOptions = createTable.withColumn(this.userColumn, DataTypes.TEXT)
.withClusteringOrder(lastClusteringColumn, ClusteringOrder.DESC)
// TODO replace w/ SchemaBuilder.unifiedCompactionStrategy() is available
.withOption("compaction", Map.of("class", "UnifiedCompactionStrategy"));
if (null != this.timeToLiveSeconds) {
createTableWithOptions = createTableWithOptions.withDefaultTimeToLiveSeconds(this.timeToLiveSeconds);
}
this.session.execute(createTableWithOptions.build());
}
}
private void ensureTableColumnsExist() {
TableMetadata tableMetadata = this.session.getMetadata()
.getKeyspace(this.schema.keyspace())
.get()
.getTable(this.schema.table())
.get();
boolean addAssistantColumn = tableMetadata.getColumn(this.assistantColumn).isEmpty();
boolean addUserColumn = tableMetadata.getColumn(this.userColumn).isEmpty();
if (addAssistantColumn || addUserColumn) {
AlterTableAddColumn alterTable = SchemaBuilder.alterTable(this.schema.keyspace(), this.schema.table());
if (addAssistantColumn) {
alterTable = alterTable.addColumn(this.assistantColumn, DataTypes.TEXT);
}
if (addUserColumn) {
alterTable = alterTable.addColumn(this.userColumn, DataTypes.TEXT);
}
SimpleStatement stmt = ((AlterTableAddColumnEnd) alterTable).build();
logger.debug("Executing {}", stmt.getQuery());
this.session.execute(stmt);
}
}
}

View File

@@ -19,6 +19,7 @@ import com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata;
import com.datastax.oss.driver.api.core.type.DataTypes;
import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.Filter.ExpressionType;
import org.springframework.ai.vectorstore.filter.Filter.Key;

View File

@@ -28,8 +28,10 @@ import com.datastax.oss.driver.api.querybuilder.delete.DeleteSelection;
import com.datastax.oss.driver.api.querybuilder.insert.InsertInto;
import com.datastax.oss.driver.api.querybuilder.insert.RegularInsert;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn;
@@ -124,6 +126,10 @@ public class CassandraVectorStore implements VectorStore, AutoCloseable {
private final Similarity similarity;
public static CassandraVectorStore create(CassandraVectorStoreConfig conf, EmbeddingModel embeddingModel) {
return new CassandraVectorStore(conf, embeddingModel);
}
public CassandraVectorStore(CassandraVectorStoreConfig conf, EmbeddingModel embeddingModel) {
Preconditions.checkArgument(null != conf, "Config must not be null");

View File

@@ -32,12 +32,14 @@ import com.datastax.oss.driver.api.querybuilder.schema.CreateTable;
import com.datastax.oss.driver.api.querybuilder.schema.CreateTableStart;
import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.lang.Nullable;
import org.springframework.ai.cassandra.SchemaUtil;
import java.net.InetSocketAddress;
import java.time.Duration;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.function.Function;
@@ -113,6 +115,8 @@ public class CassandraVectorStoreConfig implements AutoCloseable {
}
/**
* Given a string document id, return the value for each primary key column.
*
* It is a requirement that an empty {@code List<Object>} returns an example formatted
* id
*/
@@ -120,6 +124,7 @@ public class CassandraVectorStoreConfig implements AutoCloseable {
}
/** Given a list of primary key column values, return the document id. */
public interface PrimaryKeyTranslator extends Function<List<Object>, String> {
}
@@ -147,7 +152,7 @@ public class CassandraVectorStoreConfig implements AutoCloseable {
this.schema = new Schema(builder.keyspace, builder.table, builder.partitionKeys, builder.clusteringKeys,
builder.contentColumnName, builder.embeddingColumnName, builder.indexName, builder.metadataColumns);
this.disallowSchemaChanges = builder.disallowSchemaCreation;
this.disallowSchemaChanges = builder.disallowSchemaChanges;
this.returnEmbeddings = builder.returnEmbeddings;
this.documentIdTranslator = builder.documentIdTranslator;
this.primaryKeyTranslator = builder.primaryKeyTranslator;
@@ -198,7 +203,7 @@ public class CassandraVectorStoreConfig implements AutoCloseable {
private Set<SchemaColumn> metadataColumns = new HashSet<>();
private boolean disallowSchemaCreation = false;
private boolean disallowSchemaChanges = false;
private boolean returnEmbeddings = false;
@@ -254,6 +259,7 @@ public class CassandraVectorStoreConfig implements AutoCloseable {
}
public Builder withPartitionKeys(List<SchemaColumn> partitionKeys) {
Preconditions.checkArgument(!partitionKeys.isEmpty());
this.partitionKeys = partitionKeys;
return this;
}
@@ -307,7 +313,7 @@ public class CassandraVectorStoreConfig implements AutoCloseable {
}
public Builder disallowSchemaChanges() {
this.disallowSchemaCreation = true;
this.disallowSchemaChanges = true;
return this;
}
@@ -379,33 +385,17 @@ public class CassandraVectorStoreConfig implements AutoCloseable {
void ensureSchemaExists(int vectorDimension) {
if (!this.disallowSchemaChanges) {
ensureKeyspaceExists();
SchemaUtil.ensureKeyspaceExists(this.session, this.schema.keyspace);
ensureTableExists(vectorDimension);
ensureTableColumnsExist(vectorDimension);
ensureIndexesExists();
checkSchemaAgreement();
SchemaUtil.checkSchemaAgreement(session);
}
else {
checkSchemaValid(vectorDimension);
}
}
private void checkSchemaAgreement() throws IllegalStateException {
if (!this.session.checkSchemaAgreement()) {
logger.warn("Waiting for cluster schema agreement, sleeping 10s…");
try {
Thread.sleep(Duration.ofSeconds(10).toMillis());
}
catch (InterruptedException ex) {
Thread.currentThread().interrupt();
throw new IllegalStateException(ex);
}
if (!this.session.checkSchemaAgreement()) {
logger.error("no cluster schema agreement still, continuing, let's hope this works…");
}
}
}
void checkSchemaValid(int vectorDimension) {
Preconditions.checkState(this.session.getMetadata().getKeyspace(this.schema.keyspace).isPresent(),
@@ -572,16 +562,4 @@ public class CassandraVectorStoreConfig implements AutoCloseable {
}
}
private void ensureKeyspaceExists() {
if (this.session.getMetadata().getKeyspace(this.schema.keyspace).isEmpty()) {
SimpleStatement keyspaceStmt = SchemaBuilder.createKeyspace(this.schema.keyspace)
.ifNotExists()
.withSimpleStrategy(1)
.build();
logger.debug("Executing {}", keyspaceStmt.getQuery());
this.session.execute(keyspaceStmt);
}
}
}

View File

@@ -0,0 +1,91 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.memory;
import java.time.Duration;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.CassandraContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.utility.DockerImageName;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
/**
* Use `mvn failsafe:integration-test -Dit.test=CassandraChatMemoryIT`
*
* @author Mick Semb Wever
* @since 1.0.0
*/
@Testcontainers
class CassandraChatMemoryIT {
static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse("cassandra");
@Container
static CassandraContainer cassandraContainer = new CassandraContainer(DEFAULT_IMAGE_NAME.withTag("5.0"));
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(CassandraChatMemoryIT.TestApplication.class);
@Test
void ensureBeanGetsCreated() {
this.contextRunner.run(context -> {
CassandraChatMemory memory = context.getBean(CassandraChatMemory.class);
Assertions.assertNotNull(memory);
memory.conf.checkSchemaValid();
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {
@Bean
public CassandraChatMemory memory(CqlSession cqlSession) {
var conf = CassandraChatMemoryConfig.builder()
.withCqlSession(cqlSession)
.withKeyspaceName("test_" + CassandraChatMemoryConfig.DEFAULT_KEYSPACE_NAME)
.withAssistantColumnName("a")
.withUserColumnName("u")
.withTimeToLive(Duration.ofMinutes(1))
.build();
conf.dropKeyspace();
return CassandraChatMemory.create(conf);
}
@Bean
public CqlSession cqlSession() {
return new CqlSessionBuilder()
// comment next two lines out to connect to a local C* cluster
.addContactPoint(cassandraContainer.getContactPoint())
.withLocalDatacenter(cassandraContainer.getLocalDatacenter())
.build();
}
}
}

View File

@@ -195,7 +195,7 @@ class CassandraRichSchemaVectorStoreIT {
contextRunner.run(context -> {
try (CassandraVectorStore store = new CassandraVectorStore(
try (CassandraVectorStore store = CassandraVectorStore.create(
storeBuilder(context, List.of()).withFixedThreadPoolExecutorSize(nThreads).build(),
context.getBean(EmbeddingModel.class))) {

View File

@@ -396,7 +396,7 @@ class CassandraVectorStoreIT {
.build();
conf.dropKeyspace();
return new CassandraVectorStore(conf, embeddingModel);
return CassandraVectorStore.create(conf, embeddingModel);
}
@Bean
@@ -432,7 +432,7 @@ class CassandraVectorStoreIT {
CassandraVectorStoreConfig.Builder builder) {
CassandraVectorStoreConfig conf = builder.build();
conf.dropKeyspace();
return new CassandraVectorStore(conf, context.getBean(EmbeddingModel.class));
return CassandraVectorStore.create(conf, context.getBean(EmbeddingModel.class));
}
}

View File

@@ -119,7 +119,7 @@ class WikiVectorStoreExample {
})
.build();
return new CassandraVectorStore(conf, embeddingModel());
return CassandraVectorStore.create(conf, embeddingModel());
}
@Bean