Cassandra Vector Store initial impl follow up
- add concurrency to store.add(..) (bc embeddingClient is slow) - CassandraVectorStoreAutoConfiguration uses CassandraAutoConfiguration - driver profiles for production stability+performance, - small cleanups and naming fixes, - main doc tidy-up - astradb compatibility (protocol V4) – don't create embeddings again for documents that already have them similar to https://github.com/spring-projects/spring-ai/pull/413
This commit is contained in:
mck
@@ -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<Document> 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<Object> 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]
|
||||
----
|
||||
|
||||
@@ -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<InetSocketAddress> getCassandraContactPoints();
|
||||
|
||||
boolean hasCassandraLocalDatacenter();
|
||||
|
||||
String getCassandraLocalDatacenter();
|
||||
|
||||
}
|
||||
@@ -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<InetSocketAddress> 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();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
|
||||
@@ -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<Document>)} 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<Document>)} 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<Set<String>, PreparedStatement> addStmts = new HashMap<>();
|
||||
private final ConcurrentMap<Set<String>, PreparedStatement> addStmts = new ConcurrentHashMap<>();
|
||||
|
||||
private final PreparedStatement deleteStmt;
|
||||
|
||||
@@ -133,30 +153,39 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean
|
||||
|
||||
@Override
|
||||
public void add(List<Document> documents) {
|
||||
CompletableFuture[] futures = new CompletableFuture[documents.size()];
|
||||
short i = 0;
|
||||
var futures = new CompletableFuture[documents.size()];
|
||||
|
||||
int i = 0;
|
||||
for (Document d : documents) {
|
||||
List<Object> primaryKeyValues = this.conf.documentIdTranslator.apply(d.getId());
|
||||
var embedding = this.embeddingClient.embed(d).stream().map(Double::floatValue).toList();
|
||||
futures[i++] = CompletableFuture.runAsync(() -> {
|
||||
List<Object> 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<Boolean> delete(List<String> idList) {
|
||||
CompletableFuture[] futures = new CompletableFuture[idList.size()];
|
||||
short i = 0;
|
||||
int i = 0;
|
||||
for (String id : idList) {
|
||||
List<Object> 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<Document> 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<String> metadataFields) {
|
||||
if (!this.addStmts.containsKey(metadataFields)) {
|
||||
|
||||
// metadata fields that are not configured as metadata columns are not added
|
||||
Set<String> 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() {
|
||||
|
||||
@@ -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<SchemaColumn> partitionKeys, List<SchemaColumn> 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<Object> 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<SchemaColumn> 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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<Document> 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<CassandraVectorStore, CassandraVectorStoreConfig> createStore(ApplicationContext context,
|
||||
List<SchemaColumn> extraMetadataFields, boolean disallowSchemaCreation, boolean dropKeyspaceFirst)
|
||||
List<SchemaColumn> columnOverrides, boolean disallowSchemaCreation, boolean dropKeyspaceFirst)
|
||||
throws IOException {
|
||||
|
||||
Optional<SchemaColumn> 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<SchemaColumn> columnOverrides) throws IOException {
|
||||
|
||||
Optional<SchemaColumn> wikiOverride = columnOverrides.stream()
|
||||
.filter((f) -> "wiki".equals(f.name()))
|
||||
.findFirst();
|
||||
|
||||
Optional<SchemaColumn> langOverride = extraMetadataFields.stream()
|
||||
Optional<SchemaColumn> langOverride = columnOverrides.stream()
|
||||
.filter((f) -> "language".equals(f.name()))
|
||||
.findFirst();
|
||||
|
||||
Optional<SchemaColumn> titleOverride = extraMetadataFields.stream()
|
||||
Optional<SchemaColumn> titleOverride = columnOverrides.stream()
|
||||
.filter((f) -> "title".equals(f.name()))
|
||||
.findFirst();
|
||||
|
||||
Optional<SchemaColumn> chunkNoOverride = extraMetadataFields.stream()
|
||||
Optional<SchemaColumn> 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 {
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -81,25 +81,30 @@ class WikiVectorStoreExample {
|
||||
@Bean
|
||||
public CassandraVectorStore store(CqlSession cqlSession, EmbeddingClient embeddingClient) {
|
||||
|
||||
List<SchemaColumn> partitionColumns = List.of(new SchemaColumn("wiki", DataTypes.TEXT),
|
||||
new SchemaColumn("language", DataTypes.TEXT), new SchemaColumn("title", DataTypes.TEXT));
|
||||
|
||||
List<SchemaColumn> clusteringColumns = List.of(new SchemaColumn("chunk_no", DataTypes.INT),
|
||||
new SchemaColumn("bert_embedding_no", DataTypes.INT));
|
||||
|
||||
List<SchemaColumn> 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<Object> primaryKeys) -> {
|
||||
// the deliminator used to join fields together into the document's id
|
||||
// is arbitary, here "§¶" is used
|
||||
if (primaryKeys.isEmpty()) {
|
||||
return "test§¶0";
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user