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 a264c08c1..51c6110cb 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 @@ -4,9 +4,9 @@ This section walks you through setting up `CassandraVectorStore` to store docume == What is Apache Cassandra ? -link:https://cassandra.apache.org[Apache Cassandra] is a true open source distributed database reknown for scalability and high availability without compromising performance. +link:https://cassandra.apache.org[Apache Cassandra®] is a true open source distributed database reknown for linear scalability, proven fault-tolerance and low latency, making it the perfect platform for mission-critical transactional data. -Linear scalability, proven fault-tolerance and low latency on commodity hardware makes it the perfect platform for mission-critical data. Its Vector Similarity Search (VSS) is based on the JVector library that ensures best-in-class performance and relevancy. +Its Vector Similarity Search (VSS) is based on the JVector library that ensures best-in-class performance and relevancy. A vector search in Apache Cassandra is done as simply as: ``` @@ -15,9 +15,13 @@ SELECT content FROM table ORDER BY content_vector ANN OF query_embedding ; More docs on this can be read https://cassandra.apache.org/doc/latest/cassandra/getting-started/vector-search-quickstart.html[here]. -The Spring AI Cassandra Vector Store is designed to work for both brand new RAG applications as well as being able to be retrofitted on top of existing data and tables. This vector store may also equally be used for non-RAG non_AI use-cases, e.g. semantic searcing in an existing database. The Vector 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`. +This Spring AI Vector Store is designed to work for both brand new RAG applications as well as being able to be retrofitted on top of existing data and tables. -== What is JVector Vector Search ? +The store can also be used for non-RAG use-cases in an existing database, e.g. semantic searches, geo-proximity searches, etc. + +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`. + +== What is JVector ? link:https://github.com/jbellis/jvector[JVector] is a pure Java embedded vector search engine. @@ -70,13 +74,6 @@ Add these dependencies to your project: TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build file. -* If for example you want to use the OpenAI modules, remember to provide your OpenAI API Key. Set it as an environment variable like so: - -[source,bash] ----- -export SPRING_AI_OPENAI_API_KEY='Your_OpenAI_API_Key' ----- - == Usage @@ -93,21 +90,14 @@ public VectorStore vectorStore(EmbeddingClient embeddingClient) { } ---- -NOTE: It is more convenient and preferred to create the `CassandraVectorStore` as a Bean. -But if you decide you can create it manually. - [NOTE] ==== -The default configuration connects to Cassandra at localhost:9042 and will automatically create the default schema at `springframework_ai_vector.springframework_ai_vector_store`. - -Please see `CassandraVectorStoreConfig.Builder` for all the configuration options. +The default configuration connects to Cassandra at `localhost:9042` and will automatically create a default schema in keyspace `springframework`, table `ai_vector_store`. ==== [NOTE] ==== -The Cassandra Java Driver is easiest configured via the `application.conf` file on the classpath. - -More info can be found link: https://github.com/apache/cassandra-java-driver/tree/4.x/manual/core/configuration[here]. +The Cassandra Java Driver is easiest configured via an `application.conf` file on the classpath. More info https://github.com/apache/cassandra-java-driver/tree/4.x/manual/core/configuration[here]. ==== Then in your main code, create some documents: @@ -148,7 +138,7 @@ List results = vectorStore.similaritySearch( === Metadata filtering -You can leverage the generic, portable link:https://docs.spring.io/spring-ai/reference/api/vectordbs.html#_metadata_filters[metadata filters] with the CassandraVectorStore as well. Metadata fields must be configured in `CassandraVectorStoreConfig`. +You can leverage the generic, portable link:https://docs.spring.io/spring-ai/reference/api/vectordbs.html#_metadata_filters[metadata filters] with the CassandraVectorStore as well. Metadata columns must be configured in `CassandraVectorStoreConfig`. For example, you can use either the text expression language: @@ -173,7 +163,9 @@ vectorStore.similaritySearch( The portable filter expressions get automatically converted into link:https://cassandra.apache.org/doc/latest/cassandra/developing/cql/index.html[CQL queries]. -Metadata fields to be searchable need to be either primary key columns or SAI indexed. To do this configure the metadata field with the `SchemaColumnTags.INDEXED`. +For metadata columns to be searchable they must be either primary keys or SAI indexed. To make non-primary-key columns indexed configure the metadata column with the `SchemaColumnTags.INDEXED`. + + == Advanced Example: Vector Store ontop full Wikipedia dataset @@ -187,7 +179,8 @@ Create the schema in the Cassandra database first: [source,bash] ---- -wget https://raw.githubusercontent.com/datastax-labs/colbert-wikipedia-data/main/schema.cql -O colbert-wikipedia-schema.cql +wget https://s.apache.org/colbert-wikipedia-schema-cql -O colbert-wikipedia-schema.cql + cqlsh -f colbert-wikipedia-schema.cql ---- @@ -212,14 +205,14 @@ public CassandraVectorStore store(EmbeddingClient embeddingClient) { .withTableName("articles") .withPartitionKeys(partitionColumns) .withClusteringKeys(clusteringColumns) - .withContentFieldName("body") - .withEmbeddingFieldName("all_minilm_l6_v2_embedding") + .withContentColumnName("body") + .withEmbeddingColumndName("all_minilm_l6_v2_embedding") .withIndexName("all_minilm_l6_v2_ann") .disallowSchemaChanges() - .addMetadataFields(extraColumns) + .addMetadataColumns(extraColumns) .withPrimaryKeyTranslator((List primaryKeys) -> { - // the deliminator used to join fields together into the document's id - // is arbitary, here "§¶" is used + // the deliminator used to join fields together into the document's id is arbitary + // here "§¶" is used if (primaryKeys.isEmpty()) { return "test§¶0"; } @@ -243,8 +236,11 @@ public EmbeddingClient embeddingClient() { } ---- + +== Complete wikipedia dataset + And, if you would like to load the full wikipedia dataset. -First download the `simplewiki-sstable.tar` from this link https://drive.google.com/file/d/1CcMMsj8jTKRVGep4A7hmOSvaPepsaKYP/view?usp=share_link . This will take a while, the file is tens of GBs. +First download the `simplewiki-sstable.tar` from this link https://s.apache.org/simplewiki-sstable-tar . This will take a while, the file is tens of GBs. [source,bash] ---- diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraConnectionDetails.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraConnectionDetails.java deleted file mode 100644 index b67f90f6a..000000000 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraConnectionDetails.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * 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.vectorstore.cassandra; - -import java.net.InetSocketAddress; -import java.util.List; - -import org.springframework.boot.autoconfigure.service.connection.ConnectionDetails; - -/** - * @author Mick Semb Wever - * @since 1.0.0 - */ -public interface CassandraConnectionDetails extends ConnectionDetails { - - boolean hasCassandraContactPoints(); - - List getCassandraContactPoints(); - - boolean hasCassandraLocalDatacenter(); - - String getCassandraLocalDatacenter(); - -} 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 1bb93f410..f53a6ce35 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 @@ -15,16 +15,17 @@ */ package org.springframework.ai.autoconfigure.vectorstore.cassandra; -import java.net.InetSocketAddress; -import java.util.Arrays; -import java.util.List; +import java.time.Duration; -import com.google.common.base.Preconditions; +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.config.DefaultDriverOption; import org.springframework.ai.embedding.EmbeddingClient; import org.springframework.ai.vectorstore.CassandraVectorStore; import org.springframework.ai.vectorstore.CassandraVectorStoreConfig; import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.cassandra.CassandraAutoConfiguration; +import org.springframework.boot.autoconfigure.cassandra.DriverConfigLoaderBuilderCustomizer; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; import org.springframework.boot.context.properties.EnableConfigurationProperties; @@ -34,37 +35,24 @@ import org.springframework.context.annotation.Bean; * @author Mick Semb Wever * @since 1.0.0 */ -@AutoConfiguration -@ConditionalOnClass({ CassandraVectorStore.class, EmbeddingClient.class }) +@AutoConfiguration(after = CassandraAutoConfiguration.class) +@ConditionalOnClass({ CassandraVectorStore.class, EmbeddingClient.class, CqlSession.class }) @EnableConfigurationProperties(CassandraVectorStoreProperties.class) public class CassandraVectorStoreAutoConfiguration { - @Bean - @ConditionalOnMissingBean(CassandraConnectionDetails.class) - public PropertiesCassandraConnectionDetails cassandraConnectionDetails(CassandraVectorStoreProperties properties) { - return new PropertiesCassandraConnectionDetails(properties); - } - @Bean @ConditionalOnMissingBean public CassandraVectorStore vectorStore(EmbeddingClient embeddingClient, CassandraVectorStoreProperties properties, - CassandraConnectionDetails cassandraConnectionDetails) { + CqlSession cqlSession) { - var builder = CassandraVectorStoreConfig.builder(); - if (cassandraConnectionDetails.hasCassandraContactPoints()) { - for (InetSocketAddress contactPoint : cassandraConnectionDetails.getCassandraContactPoints()) { - builder = builder.addContactPoint(contactPoint); - } - } - if (cassandraConnectionDetails.hasCassandraLocalDatacenter()) { - builder = builder.withLocalDatacenter(cassandraConnectionDetails.getCassandraLocalDatacenter()); - } + var builder = CassandraVectorStoreConfig.builder().withCqlSession(cqlSession); builder = builder.withKeyspaceName(properties.getKeyspace()) .withTableName(properties.getTable()) - .withContentColumnName(properties.getContentFieldName()) - .withEmbeddingColumnName(properties.getEmbeddingFieldName()) - .withIndexName(properties.getIndexName()); + .withContentColumnName(properties.getContentColumnName()) + .withEmbeddingColumnName(properties.getEmbeddingColumnName()) + .withIndexName(properties.getIndexName()) + .withFixedThreadPoolExecutorSize(properties.getFixedThreadPoolExecutorSize()); if (properties.getDisallowSchemaCreation()) { builder = builder.disallowSchemaChanges(); @@ -73,46 +61,20 @@ public class CassandraVectorStoreAutoConfiguration { return new CassandraVectorStore(builder.build(), embeddingClient); } - private static class PropertiesCassandraConnectionDetails implements CassandraConnectionDetails { - - private final CassandraVectorStoreProperties properties; - - public PropertiesCassandraConnectionDetails(CassandraVectorStoreProperties properties) { - this.properties = properties; - } - - private String[] getCassandraContactPointHosts() { - return this.properties.getCassandraContactPointHosts().split("(,| )"); - } - - @Override - public List getCassandraContactPoints() { - - Preconditions.checkState(hasCassandraContactPoints(), "cassandraContactPointHosts has not been set"); - final int port = this.properties.getCassandraContactPointPort(); - - return Arrays.asList(getCassandraContactPointHosts()) - .stream() - .map((host) -> InetSocketAddress.createUnresolved(host, port)) - .toList(); - } - - @Override - public String getCassandraLocalDatacenter() { - Preconditions.checkState(hasCassandraLocalDatacenter(), "cassandraLocalDatacenter has not been set"); - return this.properties.getCassandraLocalDatacenter(); - } - - @Override - public boolean hasCassandraContactPoints() { - return null != this.properties.getCassandraContactPointHosts(); - } - - @Override - public boolean hasCassandraLocalDatacenter() { - return null != this.properties.getCassandraLocalDatacenter(); - } - + @Bean + public DriverConfigLoaderBuilderCustomizer driverConfigLoaderBuilderCustomizer() { + // this replaces spring-ai-cassandra-*.jar!application.conf + // as spring-boot autoconfigure will not resolve the default driver configs + return (builder) -> builder.startProfile(CassandraVectorStore.DRIVER_PROFILE_UPDATES) + .withString(DefaultDriverOption.REQUEST_CONSISTENCY, "LOCAL_QUORUM") + .withDuration(DefaultDriverOption.REQUEST_TIMEOUT, Duration.ofSeconds(1)) + .withBoolean(DefaultDriverOption.REQUEST_DEFAULT_IDEMPOTENCE, true) + .endProfile() + .startProfile(CassandraVectorStore.DRIVER_PROFILE_SEARCH) + .withString(DefaultDriverOption.REQUEST_CONSISTENCY, "LOCAL_ONE") + .withDuration(DefaultDriverOption.REQUEST_TIMEOUT, Duration.ofSeconds(10)) + .withBoolean(DefaultDriverOption.REQUEST_DEFAULT_IDEMPOTENCE, true) + .endProfile(); } } 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 73b014ab6..1f2433100 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 @@ -15,6 +15,8 @@ */ package org.springframework.ai.autoconfigure.vectorstore.cassandra; +import com.google.api.client.util.Preconditions; + import org.springframework.ai.vectorstore.CassandraVectorStoreConfig; import org.springframework.boot.context.properties.ConfigurationProperties; @@ -27,12 +29,6 @@ public class CassandraVectorStoreProperties { public static final String CONFIG_PREFIX = "spring.ai.vectorstore.cassandra"; - private String cassandraContactPointHosts = null; - - private int cassandraContactPointPort = 9042; - - private String cassandraLocalDatacenter = null; - private String keyspace = CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME; private String table = CassandraVectorStoreConfig.DEFAULT_TABLE_NAME; @@ -45,30 +41,7 @@ public class CassandraVectorStoreProperties { private boolean disallowSchemaChanges = false; - public String getCassandraContactPointHosts() { - return this.cassandraContactPointHosts; - } - - /** comma or space separated */ - public void setCassandraContactPointHosts(String cassandraContactPointHosts) { - this.cassandraContactPointHosts = cassandraContactPointHosts; - } - - public int getCassandraContactPointPort() { - return this.cassandraContactPointPort; - } - - public void setCassandraContactPointPort(int cassandraContactPointPort) { - this.cassandraContactPointPort = cassandraContactPointPort; - } - - public String getCassandraLocalDatacenter() { - return this.cassandraLocalDatacenter; - } - - public void setCassandraLocalDatacenter(String cassandraLocalDatacenter) { - this.cassandraLocalDatacenter = cassandraLocalDatacenter; - } + private int fixedThreadPoolExecutorSize = CassandraVectorStoreConfig.DEFAULT_ADD_CONCURRENCY; public String getKeyspace() { return this.keyspace; @@ -94,20 +67,20 @@ public class CassandraVectorStoreProperties { this.indexName = indexName; } - public String getContentFieldName() { + public String getContentColumnName() { return this.contentColumnName; } - public void setContentFieldName(String contentFieldName) { - this.contentColumnName = contentFieldName; + public void setContentColumnName(String contentColumnName) { + this.contentColumnName = contentColumnName; } - public String getEmbeddingFieldName() { + public String getEmbeddingColumnName() { return this.embeddingColumnName; } - public void setEmbeddingFieldName(String embeddingFieldName) { - this.embeddingColumnName = embeddingFieldName; + public void setEmbeddingColumnName(String embeddingColumnName) { + this.embeddingColumnName = embeddingColumnName; } public Boolean getDisallowSchemaCreation() { @@ -118,4 +91,13 @@ public class CassandraVectorStoreProperties { this.disallowSchemaChanges = disallowSchemaCreation; } + public int getFixedThreadPoolExecutorSize() { + return this.fixedThreadPoolExecutorSize; + } + + public void setFixedThreadPoolExecutorSize(int fixedThreadPoolExecutorSize) { + Preconditions.checkArgument(0 < fixedThreadPoolExecutorSize); + this.fixedThreadPoolExecutorSize = fixedThreadPoolExecutorSize; + } + } diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfigurationIT.java index 15530052f..119bde944 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfigurationIT.java @@ -31,6 +31,7 @@ import org.springframework.ai.transformers.TransformersEmbeddingClient; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.autoconfigure.cassandra.CassandraAutoConfiguration; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -55,18 +56,18 @@ class CassandraVectorStoreAutoConfigurationIT { ResourceUtils.getText("classpath:/test/data/great.depression.txt"), Map.of("depression", "bad"))); private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() - .withConfiguration(AutoConfigurations.of(CassandraVectorStoreAutoConfiguration.class)) + .withConfiguration( + AutoConfigurations.of(CassandraVectorStoreAutoConfiguration.class, CassandraAutoConfiguration.class)) .withUserConfiguration(Config.class) .withPropertyValues("spring.ai.vectorstore.cassandra.keyspace=test_autoconfigure") - .withPropertyValues("spring.ai.vectorstore.cassandra.contentFieldName=doc_chunk"); + .withPropertyValues("spring.ai.vectorstore.cassandra.contentColumnName=doc_chunk"); @Test void addAndSearch() { - contextRunner - .withPropertyValues("spring.ai.vectorstore.cassandra.cassandraContactPointHosts=" + getContactPointHost()) - .withPropertyValues("spring.ai.vectorstore.cassandra.cassandraContactPointPort=" + getContactPointPort()) - .withPropertyValues("spring.ai.vectorstore.cassandra.cassandraLocalDatacenter=" - + cassandraContainer.getLocalDatacenter()) + contextRunner.withPropertyValues("spring.cassandra.contactPoints=" + getContactPointHost()) + .withPropertyValues("spring.cassandra.port=" + getContactPointPort()) + .withPropertyValues("spring.cassandra.localDatacenter=" + cassandraContainer.getLocalDatacenter()) + .withPropertyValues("spring.ai.vectorstore.cassandra.fixedThreadPoolExecutorSize=8") .run(context -> { VectorStore vectorStore = context.getBean(VectorStore.class); diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStorePropertiesTests.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStorePropertiesTests.java index c5a3e4d0f..ca5c678d7 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStorePropertiesTests.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStorePropertiesTests.java @@ -30,39 +30,34 @@ class CassandraVectorStorePropertiesTests { @Test void defaultValues() { var props = new CassandraVectorStoreProperties(); - assertThat(props.getCassandraContactPointHosts()).isNull(); - assertThat(props.getCassandraContactPointPort()).isEqualTo(9042); - assertThat(props.getCassandraLocalDatacenter()).isNull(); assertThat(props.getKeyspace()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME); assertThat(props.getTable()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_TABLE_NAME); - assertThat(props.getContentFieldName()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_CONTENT_COLUMN_NAME); - assertThat(props.getEmbeddingFieldName()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_EMBEDDING_COLUMN_NAME); + assertThat(props.getContentColumnName()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_CONTENT_COLUMN_NAME); + assertThat(props.getEmbeddingColumnName()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_EMBEDDING_COLUMN_NAME); assertThat(props.getIndexName()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_INDEX_NAME); assertThat(props.getDisallowSchemaCreation()).isFalse(); + assertThat(props.getFixedThreadPoolExecutorSize()) + .isEqualTo(CassandraVectorStoreConfig.DEFAULT_ADD_CONCURRENCY); } @Test void customValues() { var props = new CassandraVectorStoreProperties(); - props.setCassandraContactPointHosts("127.0.0.1,127.0.0.2"); - props.setCassandraContactPointPort(9043); - props.setCassandraLocalDatacenter("dc1"); props.setKeyspace("my_keyspace"); props.setTable("my_table"); - props.setContentFieldName("my_content"); - props.setEmbeddingFieldName("my_vector"); + props.setContentColumnName("my_content"); + props.setEmbeddingColumnName("my_vector"); props.setIndexName("my_sai"); props.setDisallowSchemaCreation(true); + props.setFixedThreadPoolExecutorSize(10); - assertThat(props.getCassandraContactPointHosts()).isEqualTo("127.0.0.1,127.0.0.2"); - assertThat(props.getCassandraContactPointPort()).isEqualTo(9043); - assertThat(props.getCassandraLocalDatacenter()).isEqualTo("dc1"); assertThat(props.getKeyspace()).isEqualTo("my_keyspace"); assertThat(props.getTable()).isEqualTo("my_table"); - assertThat(props.getContentFieldName()).isEqualTo("my_content"); - assertThat(props.getEmbeddingFieldName()).isEqualTo("my_vector"); + assertThat(props.getContentColumnName()).isEqualTo("my_content"); + assertThat(props.getEmbeddingColumnName()).isEqualTo("my_vector"); assertThat(props.getIndexName()).isEqualTo("my_sai"); assertThat(props.getDisallowSchemaCreation()).isTrue(); + assertThat(props.getFixedThreadPoolExecutorSize()).isEqualTo(10); } } diff --git a/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java index 4a8b681a5..3efd44034 100644 --- a/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java +++ b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java @@ -79,7 +79,7 @@ final class CassandraFilterExpressionConverter extends AbstractFilterExpressionC // TODO SAI supports collections // reach out to mck@apache.org if you'd like these implemented // case CONTAINS -> context.append(" CONTAINS "); - // case CONTAINS_KEY -> context.append(" CONTAINS KEY "); + // case CONTAINS_KEY -> context.append(" CONTAINS_KEY "); default -> throw new UnsupportedOperationException( String.format("Expression type %s not yet implemented. Patches welcome.", type)); } diff --git a/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java index 6f651c944..4e2732d58 100644 --- a/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java +++ b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java @@ -17,16 +17,20 @@ package org.springframework.ai.vectorstore; import java.util.ArrayList; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import com.datastax.oss.driver.api.core.cql.BoundStatement; 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.core.cql.SimpleStatement; import com.datastax.oss.driver.api.core.data.CqlVector; import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata; import com.datastax.oss.driver.api.querybuilder.QueryBuilder; @@ -40,6 +44,7 @@ import org.slf4j.LoggerFactory; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.vectorstore.CassandraVectorStoreConfig; import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; import org.springframework.beans.factory.InitializingBean; @@ -53,13 +58,13 @@ import org.springframework.beans.factory.InitializingBean; * fields in the documents to be stored alongside the vector and content data. * * This class requires a CassandraVectorStoreConfig configuration object for - * initialization, which includes settings like connection details, index name, field + * initialization, which includes settings like connection details, index name, column * names, etc. It also requires an EmbeddingClient to convert documents into embeddings * before storing them. * * A schema matching the configuration is automatically created if it doesn't exist. * Missing columns and indexes in existing tables will also be automatically created. - * Disable this with the disallowSchemaCreation. + * Disable this with the CassandraVectorStoreConfig#disallowSchemaChanges(). * * This class is designed to work with brand new tables that it creates for you, or on top * of existing Cassandra tables. The latter is appropriate when wanting to keep data in @@ -69,9 +74,20 @@ import org.springframework.beans.factory.InitializingBean; * Instances of this class are not dynamic against server-side schema changes. If you * change the schema server-side you need a new CassandraVectorStore instance. * + * When adding documents with the method {@link #add(List)} it first calls + * embeddingClient to create the embeddings. This is slow. Configure + * {@link CassandraVectorStoreConfig.Builder#withFixedThreadPoolExecutorSize(int)} + * accordingly to improve performance so embeddings are created and the documents are + * added concurrently. The default concurrency is 16 + * ({@link CassandraVectorStoreConfig#DEFAULT_ADD_CONCURRENCY}). Remote transformers + * probably want higher concurrency, and local transformers may need lower concurrency. + * This concurrency limit does not need to be higher than the max parallel calls made to + * the {@link #add(List)} method multiplied by the list size. This setting can + * also serve as a protecting throttle against your embedding model. + * * @author Mick Semb Wever * @see VectorStore - * @see CassandraVectorStoreConfig + * @see org.springframework.ai.vectorstore.CassandraVectorStoreConfig * @see EmbeddingClient * @since 1.0.0 */ @@ -87,10 +103,14 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean } - private static final String QUERY_FORMAT = "select %s,%s,%s%s from %s.%s ? order by %s ann of ? limit ?"; - public static final String SIMILARITY_FIELD_NAME = "similarity_score"; + public static final String DRIVER_PROFILE_UPDATES = "spring-ai-updates"; + + public static final String DRIVER_PROFILE_SEARCH = "spring-ai-search"; + + private static final String QUERY_FORMAT = "select %s,%s,%s%s from %s.%s ? order by %s ann of ? limit ?"; + private static final Logger logger = LoggerFactory.getLogger(CassandraVectorStore.class); private final CassandraVectorStoreConfig conf; @@ -99,7 +119,7 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean private final FilterExpressionConverter filterExpressionConverter; - private final Map, PreparedStatement> addStmts = new HashMap<>(); + private final ConcurrentMap, PreparedStatement> addStmts = new ConcurrentHashMap<>(); private final PreparedStatement deleteStmt; @@ -133,30 +153,39 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean @Override public void add(List documents) { - CompletableFuture[] futures = new CompletableFuture[documents.size()]; - short i = 0; + var futures = new CompletableFuture[documents.size()]; + + int i = 0; for (Document d : documents) { - List primaryKeyValues = this.conf.documentIdTranslator.apply(d.getId()); - var embedding = this.embeddingClient.embed(d).stream().map(Double::floatValue).toList(); + futures[i++] = CompletableFuture.runAsync(() -> { + List primaryKeyValues = this.conf.documentIdTranslator.apply(d.getId()); - BoundStatementBuilder builder = prepareAddStatement(d.getMetadata().keySet()).boundStatementBuilder(); - for (int k = 0; k < primaryKeyValues.size(); ++k) { - SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k); - builder = builder.set(keyColumn.name(), primaryKeyValues.get(k), keyColumn.javaType()); - } + var embedding = (null != d.getEmbedding() && !d.getEmbedding().isEmpty() ? d.getEmbedding() + : this.embeddingClient.embed(d)) + .stream() + .map(Double::floatValue) + .toList(); - builder = builder.setString(this.conf.schema.content(), d.getContent()) - .setVector(this.conf.schema.embedding(), CqlVector.newInstance(embedding), Float.class); + BoundStatementBuilder builder = prepareAddStatement(d.getMetadata().keySet()).boundStatementBuilder(); + for (int k = 0; k < primaryKeyValues.size(); ++k) { + SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k); + builder = builder.set(keyColumn.name(), primaryKeyValues.get(k), keyColumn.javaType()); + } - for (var metadataColumn : this.conf.schema.metadataColumns() - .stream() - .filter((mc) -> d.getMetadata().containsKey(mc.name())) - .toList()) { + builder = builder.setString(this.conf.schema.content(), d.getContent()) + .setVector(this.conf.schema.embedding(), CqlVector.newInstance(embedding), Float.class); - builder = builder.set(metadataColumn.name(), d.getMetadata().get(metadataColumn.name()), - metadataColumn.javaType()); - } - futures[i++] = this.conf.session.executeAsync(builder.build()).toCompletableFuture(); + for (var metadataColumn : this.conf.schema.metadataColumns() + .stream() + .filter((mc) -> d.getMetadata().containsKey(mc.name())) + .toList()) { + + builder = builder.set(metadataColumn.name(), d.getMetadata().get(metadataColumn.name()), + metadataColumn.javaType()); + } + BoundStatement s = builder.build().setExecutionProfileName(DRIVER_PROFILE_UPDATES); + this.conf.session.execute(s); + }, this.conf.executor); } CompletableFuture.allOf(futures).join(); } @@ -164,7 +193,7 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean @Override public Optional delete(List idList) { CompletableFuture[] futures = new CompletableFuture[idList.size()]; - short i = 0; + int i = 0; for (String id : idList) { List primaryKeyValues = this.conf.documentIdTranslator.apply(id); BoundStatement s = this.deleteStmt.bind(primaryKeyValues.toArray()); @@ -191,8 +220,9 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean String query = String.format(this.similarityStmt, cqlVector, whereClause, cqlVector, request.getTopK()); List documents = new ArrayList<>(); logger.trace("Executing {}", query); + SimpleStatement s = SimpleStatement.newInstance(query).setExecutionProfileName(DRIVER_PROFILE_SEARCH); - for (Row row : this.conf.session.execute(query)) { + for (Row row : this.conf.session.execute(s)) { float score = row.getFloat(0); if (score < request.getSimilarityThreshold()) { break; @@ -248,7 +278,16 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean } private PreparedStatement prepareAddStatement(Set metadataFields) { - if (!this.addStmts.containsKey(metadataFields)) { + + // metadata fields that are not configured as metadata columns are not added + Set fieldsThatAreColumns = new HashSet<>(this.conf.schema.metadataColumns() + .stream() + .map((mc) -> mc.name()) + .filter((mc) -> metadataFields.contains(mc)) + .toList()); + + return this.addStmts.computeIfAbsent(fieldsThatAreColumns, (fields) -> { + RegularInsert stmt = null; InsertInto stmtStart = QueryBuilder.insertInto(this.conf.schema.keyspace(), this.conf.schema.table()); @@ -262,17 +301,11 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean stmt = stmt.value(this.conf.schema.content(), QueryBuilder.bindMarker(this.conf.schema.content())) .value(this.conf.schema.embedding(), QueryBuilder.bindMarker(this.conf.schema.embedding())); - for (String metadataField : this.conf.schema.metadataColumns() - .stream() - .map((mc) -> mc.name()) - .filter((mc) -> metadataFields.contains(mc)) - .toList()) { - + for (String metadataField : fields) { stmt = stmt.value(metadataField, QueryBuilder.bindMarker(metadataField)); } - this.addStmts.putIfAbsent(metadataFields, this.conf.session.prepare(stmt.build())); - } - return this.addStmts.get(metadataFields); + return this.conf.session.prepare(stmt.build()); + }); } private String similaritySearchStatement() { diff --git a/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java index d85d1ee4b..38b39e16f 100644 --- a/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java +++ b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java @@ -22,6 +22,8 @@ import java.util.HashSet; import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.concurrent.Executor; +import java.util.concurrent.Executors; import java.util.function.Function; import java.util.stream.Stream; @@ -48,15 +50,17 @@ import org.slf4j.LoggerFactory; /** * Configuration for the Cassandra vector store. * - * All metadata fields configured to the store will be fetched and added to all queried + * All metadata columns configured to the store will be fetched and added to all queried * documents. * - * If you wish to metadata search against a field its 'searchable' argument must be true. + * To filter expression search against a metadata column configure it with + * SchemaColumnTags.INDEXED * * The Cassandra Java Driver is configured via the application.conf resource found in the * classpath. See * https://github.com/apache/cassandra-java-driver/tree/4.x/manual/core/configuration * + * @author Mick Semb Wever * @since 1.0.0 */ public final class CassandraVectorStoreConfig implements AutoCloseable { @@ -73,6 +77,8 @@ public final class CassandraVectorStoreConfig implements AutoCloseable { public static final String DEFAULT_EMBEDDING_COLUMN_NAME = "embedding"; + public static final int DEFAULT_ADD_CONCURRENCY = 16; + private static final Logger logger = LoggerFactory.getLogger(CassandraVectorStore.class); record Schema(String keyspace, String table, List partitionKeys, List clusteringKeys, @@ -127,6 +133,8 @@ public final class CassandraVectorStoreConfig implements AutoCloseable { final PrimaryKeyTranslator primaryKeyTranslator; + final Executor executor; + private final boolean closeSessionOnClose; private CassandraVectorStoreConfig(Builder builder) { @@ -139,6 +147,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable { this.disallowSchemaChanges = builder.disallowSchemaCreation; this.documentIdTranslator = builder.documentIdTranslator; this.primaryKeyTranslator = builder.primaryKeyTranslator; + this.executor = Executors.newFixedThreadPool(builder.fixedThreadPoolExecutorSize); } public static Builder builder() { @@ -187,6 +196,8 @@ public final class CassandraVectorStoreConfig implements AutoCloseable { private boolean disallowSchemaCreation = false; + private int fixedThreadPoolExecutorSize = DEFAULT_ADD_CONCURRENCY; + private DocumentIdTranslator documentIdTranslator = (String id) -> List.of(id); private PrimaryKeyTranslator primaryKeyTranslator = (List primaryKeyColumns) -> { @@ -261,20 +272,27 @@ public final class CassandraVectorStoreConfig implements AutoCloseable { return this; } - public Builder addMetadataColumn(SchemaColumn... fields) { + public Builder addMetadataColumns(SchemaColumn... columns) { Builder builder = this; - for (SchemaColumn f : fields) { + for (SchemaColumn f : columns) { builder = builder.addMetadataColumn(f); } return builder; } - public Builder addMetadataColumn(SchemaColumn field) { + public Builder addMetadataColumns(List columns) { + Builder builder = this; + this.metadataColumns.addAll(columns); + return builder; + } - Preconditions.checkArgument(this.metadataColumns.stream().noneMatch((sc) -> sc.name().equals(field.name())), - "A metadata field with name %s has already been added", field.name()); + public Builder addMetadataColumn(SchemaColumn column) { - this.metadataColumns.add(field); + Preconditions.checkArgument( + this.metadataColumns.stream().noneMatch((sc) -> sc.name().equals(column.name())), + "A metadata column with name %s has already been added", column.name()); + + this.metadataColumns.add(column); return this; } @@ -283,6 +301,18 @@ public final class CassandraVectorStoreConfig implements AutoCloseable { return this; } + /** + * Executor to use when adding documents. The hotspot is the call to the + * embeddingClient. For remote transformers you probably want a higher value to + * utilize network. For local transformers you probably want a lower value to + * avoid saturation. + **/ + public Builder withFixedThreadPoolExecutorSize(int threads) { + Preconditions.checkArgument(0 < threads); + this.fixedThreadPoolExecutorSize = threads; + return this; + } + public Builder withDocumentIdTranslator(DocumentIdTranslator documentIdTranslator) { this.documentIdTranslator = documentIdTranslator; return this; @@ -480,7 +510,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable { if (column.isPresent()) { Preconditions.checkArgument(column.get().getType().equals(metadata.type()), - "Cannot change type on metadata field %s from %s to %s", metadata.name(), + "Cannot change type on metadata column %s from %s to %s", metadata.name(), column.get().getType(), metadata.type()); } else { diff --git a/vector-stores/spring-ai-cassandra/src/main/resources/application.conf b/vector-stores/spring-ai-cassandra/src/main/resources/application.conf new file mode 100644 index 000000000..91b1c800e --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/main/resources/application.conf @@ -0,0 +1,24 @@ +# Reference configuration for the DataStax Java driver for Apache Cassandra® +# see https://github.com/apache/cassandra-java-driver/tree/4.x/manual/core/configuration +# +# +# when using spring-boot autoconfigure this will not be used +# instead CassandraVectorStoreAutoConfiguration.driverConfigLoaderBuilderCustomizer() is used +datastax-java-driver { + profiles { + spring-ai-updates { + basic.request { + consistency = LOCAL_QUORUM + timeout = 1 seconds + default-idempotence = true + } + } + spring-ai-search { + basic.request { + consistency = LOCAL_ONE + timeout = 10 seconds + default-idempotence = true + } + } + } +} \ No newline at end of file diff --git a/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java index 91ec53e50..2dc59afae 100644 --- a/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java +++ b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java @@ -17,23 +17,29 @@ package org.springframework.ai.vectorstore; import java.io.IOException; import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Optional; -import org.junit.jupiter.api.Assertions; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executors; +import java.util.concurrent.ThreadLocalRandom; import com.datastax.oss.driver.api.core.CqlSession; import com.datastax.oss.driver.api.core.CqlSessionBuilder; import com.datastax.oss.driver.api.core.servererrors.InvalidQueryException; import com.datastax.oss.driver.api.core.servererrors.SyntaxError; import com.datastax.oss.driver.api.core.type.DataTypes; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.testcontainers.containers.CassandraContainer; import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.shaded.org.apache.commons.lang3.RandomStringUtils; import org.testcontainers.utility.DockerImageName; import org.springframework.ai.document.Document; @@ -174,6 +180,51 @@ class CassandraRichSchemaVectorStoreIT { }); } + @Test + void addAndSearchPoormansBench() { + // todo – replace with JMH (parameters: nThreads, rounds, runs, docsPerAdd) + int nThreads = CassandraVectorStoreConfig.DEFAULT_ADD_CONCURRENCY; + int runs = 10; // 100; + int docsPerAdd = 12; // 128; + int rounds = 3; + + contextRunner.run(context -> { + + try (CassandraVectorStore store = new CassandraVectorStore( + storeBuilder(context, List.of()).withFixedThreadPoolExecutorSize(nThreads).build(), + context.getBean(EmbeddingClient.class))) { + + var executor = Executors.newFixedThreadPool((int) (nThreads * 1.2)); + for (int k = 0; k < rounds; ++k) { + long start = System.nanoTime(); + var futures = new CompletableFuture[runs]; + for (int j = 0; j < runs; ++j) { + futures[j] = CompletableFuture.runAsync(() -> { + List documents = new ArrayList<>(); + for (int i = docsPerAdd; i >= 0; --i) { + + documents.add(new Document( + RandomStringUtils.randomAlphanumeric(4) + "§¶" + + ThreadLocalRandom.current().nextInt(1, 10), + RandomStringUtils.randomAlphanumeric(1024), Map.of("revision", + ThreadLocalRandom.current().nextInt(1, 100000), "id", 1000))); + } + store.add(documents); + + var results = store.similaritySearch( + SearchRequest.query(RandomStringUtils.randomAlphanumeric(20)).withTopK(10)); + + assertThat(results).hasSize(10); + }, executor); + } + CompletableFuture.allOf(futures).join(); + long time = System.nanoTime() - start; + logger.info("add+search took an average of {} ms", Duration.ofNanos(time / runs).toMillis()); + } + } + }); + } + @Test void searchWithPartitionFilter() throws InterruptedException { contextRunner.run(context -> { @@ -456,22 +507,37 @@ class CassandraRichSchemaVectorStoreIT { } private StoreWrapper createStore(ApplicationContext context, - List extraMetadataFields, boolean disallowSchemaCreation, boolean dropKeyspaceFirst) + List columnOverrides, boolean disallowSchemaCreation, boolean dropKeyspaceFirst) throws IOException { - Optional wikiOverride = extraMetadataFields.stream() + CassandraVectorStoreConfig.Builder builder = storeBuilder(context, columnOverrides); + if (disallowSchemaCreation) { + builder = builder.disallowSchemaChanges(); + } + + CassandraVectorStoreConfig conf = builder.build(); + if (dropKeyspaceFirst) { + conf.dropKeyspace(); + } + return new StoreWrapper(new CassandraVectorStore(conf, context.getBean(EmbeddingClient.class)), conf); + } + + static CassandraVectorStoreConfig.Builder storeBuilder(ApplicationContext context, + List columnOverrides) throws IOException { + + Optional wikiOverride = columnOverrides.stream() .filter((f) -> "wiki".equals(f.name())) .findFirst(); - Optional langOverride = extraMetadataFields.stream() + Optional langOverride = columnOverrides.stream() .filter((f) -> "language".equals(f.name())) .findFirst(); - Optional titleOverride = extraMetadataFields.stream() + Optional titleOverride = columnOverrides.stream() .filter((f) -> "title".equals(f.name())) .findFirst(); - Optional chunkNoOverride = extraMetadataFields.stream() + Optional chunkNoOverride = columnOverrides.stream() .filter((f) -> "chunk_no".equals(f.name())) .findFirst(); @@ -493,7 +559,7 @@ class CassandraRichSchemaVectorStoreIT { .withEmbeddingColumnName("all_minilm_l6_v2_embedding") .withIndexName("all_minilm_l6_v2_ann") - .addMetadataColumn(new SchemaColumn("revision", DataTypes.INT), + .addMetadataColumns(new SchemaColumn("revision", DataTypes.INT), new SchemaColumn("id", DataTypes.INT, CassandraVectorStoreConfig.SchemaColumnTags.INDEXED)) // this store uses '§¶' as a deliminator in the document id between db columns @@ -511,21 +577,7 @@ class CassandraRichSchemaVectorStoreIT { return List.of("simplewiki", "en", title, chunk_no); }); - for (SchemaColumn cf : extraMetadataFields) { - if (!partitionKeys.contains(cf) && !clusteringKeys.contains(cf)) { - builder = builder.addMetadataColumn(cf); - } - } - - if (disallowSchemaCreation) { - builder = builder.disallowSchemaChanges(); - } - - CassandraVectorStoreConfig conf = builder.build(); - if (dropKeyspaceFirst) { - conf.dropKeyspace(); - } - return new StoreWrapper(new CassandraVectorStore(conf, context.getBean(EmbeddingClient.class)), conf); + return builder; } private void executeCqlFile(ApplicationContext context, String filename) throws IOException { diff --git a/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java index 0984a067f..d1fac3901 100644 --- a/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java +++ b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java @@ -345,8 +345,9 @@ class CassandraVectorStoreIT { public CassandraVectorStore store(CqlSession cqlSession, EmbeddingClient embeddingClient) { CassandraVectorStoreConfig conf = storeBuilder(cqlSession) - .addMetadataColumn(new SchemaColumn("meta1", DataTypes.TEXT), new SchemaColumn("meta2", DataTypes.TEXT), - new SchemaColumn("country", DataTypes.TEXT), new SchemaColumn("year", DataTypes.SMALLINT)) + .addMetadataColumns(new SchemaColumn("meta1", DataTypes.TEXT), + new SchemaColumn("meta2", DataTypes.TEXT), new SchemaColumn("country", DataTypes.TEXT), + new SchemaColumn("year", DataTypes.SMALLINT)) .build(); conf.dropKeyspace(); @@ -378,7 +379,7 @@ class CassandraVectorStoreIT { private CassandraVectorStore createTestStore(ApplicationContext context, SchemaColumn... metadataFields) { CassandraVectorStoreConfig.Builder builder = storeBuilder(context.getBean(CqlSession.class)) - .addMetadataColumn(metadataFields); + .addMetadataColumns(metadataFields); CassandraVectorStoreConfig conf = builder.build(); conf.dropKeyspace(); diff --git a/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java index 76028e650..910dac85a 100644 --- a/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java +++ b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java @@ -81,25 +81,30 @@ class WikiVectorStoreExample { @Bean public CassandraVectorStore store(CqlSession cqlSession, EmbeddingClient embeddingClient) { + List partitionColumns = List.of(new SchemaColumn("wiki", DataTypes.TEXT), + new SchemaColumn("language", DataTypes.TEXT), new SchemaColumn("title", DataTypes.TEXT)); + + List clusteringColumns = List.of(new SchemaColumn("chunk_no", DataTypes.INT), + new SchemaColumn("bert_embedding_no", DataTypes.INT)); + + List extraColumns = List.of(new SchemaColumn("revision", DataTypes.INT), + new SchemaColumn("id", DataTypes.INT)); + CassandraVectorStoreConfig conf = CassandraVectorStoreConfig.builder() .withCqlSession(cqlSession) .withKeyspaceName("wikidata") .withTableName("articles") - - .withPartitionKeys(List.of(new SchemaColumn("wiki", DataTypes.TEXT), - new SchemaColumn("language", DataTypes.TEXT), new SchemaColumn("title", DataTypes.TEXT))) - - .withClusteringKeys(List.of(new SchemaColumn("chunk_no", DataTypes.INT), - new SchemaColumn("bert_embedding_no", DataTypes.INT))) - + .withPartitionKeys(partitionColumns) + .withClusteringKeys(clusteringColumns) .withContentColumnName("body") .withEmbeddingColumnName("all_minilm_l6_v2_embedding") .withIndexName("all_minilm_l6_v2_ann") .disallowSchemaChanges() - - .addMetadataColumn(new SchemaColumn("revision", DataTypes.INT), new SchemaColumn("id", DataTypes.INT)) + .addMetadataColumns(extraColumns) .withPrimaryKeyTranslator((List primaryKeys) -> { + // the deliminator used to join fields together into the document's id + // is arbitary, here "§¶" is used if (primaryKeys.isEmpty()) { return "test§¶0"; }