diff --git a/.gitignore b/.gitignore index c2fb8fc27..19db9000b 100644 --- a/.gitignore +++ b/.gitignore @@ -38,4 +38,6 @@ package.json shell.log -.profiler \ No newline at end of file +.profiler +/spring-ai-spring-boot-autoconfigure/nbproject/ +/vector-stores/spring-ai-cassandra-store/nbproject/ diff --git a/pom.xml b/pom.xml index ce1bc39a4..84744dc02 100644 --- a/pom.xml +++ b/pom.xml @@ -170,7 +170,7 @@ 3.25.2 - 4.18.0 + 4.18.1 3.0.1 0.1.4 2.20.11 diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/ChatMemory.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/ChatMemory.java index 9e1cad6f3..6d383c0c4 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/ChatMemory.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/memory/ChatMemory.java @@ -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 { diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc index eb2887a46..fa494278c 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc @@ -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 diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/apache-cassandra.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/apache-cassandra.adoc index 6e1b407e1..e064f2649 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/apache-cassandra.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/apache-cassandra.adoc @@ -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. diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/CommonChatMemoryProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/CommonChatMemoryProperties.java new file mode 100644 index 000000000..a635c70f6 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/CommonChatMemoryProperties.java @@ -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; + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryAutoConfiguration.java new file mode 100644 index 000000000..b84011dfe --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryAutoConfiguration.java @@ -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()); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryProperties.java new file mode 100644 index 000000000..91d2252bc --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryProperties.java @@ -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; + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/CommonVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/CommonVectorStoreProperties.java similarity index 83% rename from spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/CommonVectorStoreProperties.java rename to spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/CommonVectorStoreProperties.java index 8dcf45b23..cf429fe1a 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/CommonVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/CommonVectorStoreProperties.java @@ -1,4 +1,4 @@ -package org.springframework.ai.autoconfigure; +package org.springframework.ai.autoconfigure.vectorstore; /** * @author Josh Long diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/azure/AzureVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/azure/AzureVectorStoreProperties.java index d2ad4a8ea..6de3a5730 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/azure/AzureVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/azure/AzureVectorStoreProperties.java @@ -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; diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfiguration.java index eb053338f..536ae5154 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfiguration.java @@ -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 diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreProperties.java index 9725e7559..f84a2f0e4 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreProperties.java @@ -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() { diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreProperties.java index 8283cb6ac..1768ed81d 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreProperties.java @@ -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; diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java index cf456b137..ba0100bb3 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java @@ -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; diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreProperties.java index 15a0b0cce..2a4b82864 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreProperties.java @@ -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; diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreProperties.java index 12c16fdb8..799b58261 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreProperties.java @@ -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; /** diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java index 851cd100a..5d782b8a9 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java @@ -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; diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/oracle/OracleVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/oracle/OracleVectorStoreProperties.java index 2d70e23f7..2d5eb2f7f 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/oracle/OracleVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/oracle/OracleVectorStoreProperties.java @@ -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; diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreProperties.java index 783041c76..b45541746 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreProperties.java @@ -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; diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/qdrant/QdrantVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/qdrant/QdrantVectorStoreProperties.java index 795be25b8..880f6925d 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/qdrant/QdrantVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/qdrant/QdrantVectorStoreProperties.java @@ -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; diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/redis/RedisVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/redis/RedisVectorStoreProperties.java index 8edff4ba0..eb79a1a40 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/redis/RedisVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/redis/RedisVectorStoreProperties.java @@ -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; /** diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/weaviate/WeaviateVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/weaviate/WeaviateVectorStoreProperties.java index 654a22409..7ff377963 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/weaviate/WeaviateVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/weaviate/WeaviateVectorStoreProperties.java @@ -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; diff --git a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports index c6c8d7a28..bcc472eb6 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports +++ b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -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 \ No newline at end of file +org.springframework.ai.autoconfigure.vertexai.embedding.VertexAiEmbeddingAutoConfiguration +org.springframework.ai.autoconfigure.chat.memory.cassandra.CassandraAutoConfiguration \ No newline at end of file diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryAutoConfigurationIT.java new file mode 100644 index 000000000..86df7f40c --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryAutoConfigurationIT.java @@ -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()); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryPropertiesTest.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryPropertiesTest.java new file mode 100644 index 000000000..e3df66722 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryPropertiesTest.java @@ -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(); + } + +} diff --git a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/cassandra/SchemaUtil.java b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/cassandra/SchemaUtil.java new file mode 100644 index 000000000..f945b9db8 --- /dev/null +++ b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/cassandra/SchemaUtil.java @@ -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); + } + } + +} diff --git a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/CassandraChatMemory.java b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/CassandraChatMemory.java new file mode 100644 index 000000000..7dc82f370 --- /dev/null +++ b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/CassandraChatMemory.java @@ -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 +CassandraChatMemory.create(CassandraChatMemoryConfig.builder().withTimeToLive(Duration.ofDays(1)).build()); + + * + * 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 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 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 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 get(String sessionId, int lastN) { + + List 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 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()); + } + +} diff --git a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/CassandraChatMemoryConfig.java b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/CassandraChatMemoryConfig.java new file mode 100644 index 000000000..ab046e2ec --- /dev/null +++ b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/CassandraChatMemoryConfig.java @@ -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 partitionKeys, List clusteringKeys) { + } + + public record SchemaColumn(String name, DataType type) { + + public GenericType javaType() { + return CodecRegistry.DEFAULT.codecFor(type).getJavaType(); + } + } + + /** Given a string sessionId, return the value for each primary key column. */ + public interface SessionIdToPrimaryKeysTranslator extends Function> { + + } + + 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 partitionKeys = List.of(new SchemaColumn(DEFAULT_SESSION_ID_NAME, DataTypes.TEXT)); + + private List 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 partitionKeys) { + Preconditions.checkArgument(!partitionKeys.isEmpty()); + this.partitionKeys = partitionKeys; + return this; + } + + public Builder withClusteringKeys(List 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); + } + } + +} diff --git a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java index f3270ac81..aef2a56b2 100644 --- a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java +++ b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java @@ -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; diff --git a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java index 2f4e138eb..09cce8a67 100644 --- a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java +++ b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java @@ -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"); diff --git a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java index e102c3091..e50e8008c 100644 --- a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java +++ b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java @@ -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} 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, 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 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 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); - } - } - } diff --git a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/chat/memory/CassandraChatMemoryIT.java b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/chat/memory/CassandraChatMemoryIT.java new file mode 100644 index 000000000..eb824e40a --- /dev/null +++ b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/chat/memory/CassandraChatMemoryIT.java @@ -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(); + } + + } + +} diff --git a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java index 26b336f36..ad9dfff25 100644 --- a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java +++ b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java @@ -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))) { diff --git a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java index b362cc2c4..4f1088697 100644 --- a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java +++ b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java @@ -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)); } } diff --git a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java index 301b61c49..c772df762 100644 --- a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java +++ b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java @@ -119,7 +119,7 @@ class WikiVectorStoreExample { }) .build(); - return new CassandraVectorStore(conf, embeddingModel()); + return CassandraVectorStore.create(conf, embeddingModel()); } @Bean