Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add builder pattern to CassandraVectorStore and refactor package name

Browse Source 
Introduces a builder pattern for configuring CassandraVectorStore instances and
moves the implementation to the org.springframework.ai.vectorstore.cassandra
package. This change:

- Makes configuration more flexible and type-safe through builder methods
- Improves code organization by moving to a dedicated vector store package
- Deprecates old constructors in favor of the builder pattern
- Adds comprehensive validation of configuration options
- Enables better IDE support through method chaining
- The builder pattern provides a more maintainable and user-friendly way to
  configure vector stores while ensuring configuration validity at compile time.
This commit is contained in:
Soby Chacko
2024-12-11 19:08:08 -05:00
committed by Mark Pollack
parent d3d34c9215
commit 25123a5364
21 changed files with 1447 additions and 702 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/CassandraChatMemory.java → vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/cassandra/CassandraChatMemory.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.chat.memory;
package org.springframework.ai.chat.memory.cassandra;
import java.time.Instant;
import java.util.ArrayList;
@@ -32,7 +32,8 @@ import com.datastax.oss.driver.api.querybuilder.insert.RegularInsert;
import com.datastax.oss.driver.api.querybuilder.select.Select;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import org.springframework.ai.chat.memory.CassandraChatMemoryConfig.SchemaColumn;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryConfig.SchemaColumn;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
@@ -42,7 +43,7 @@ import org.springframework.ai.chat.messages.UserMessage;
CassandraChatMemory.create(CassandraChatMemoryConfig.builder().withTimeToLive(Duration.ofDays(1)).build());
</code>
*
* For example @see org.springframework.ai.chat.memory.CassandraChatMemory
* For example @see org.springframework.ai.chat.memory.cassandra.CassandraChatMemory
*
* @author Mick Semb Wever
* @since 1.0.0

2
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/CassandraChatMemoryConfig.java → vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/cassandra/CassandraChatMemoryConfig.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.chat.memory;
package org.springframework.ai.chat.memory.cassandra;
import java.net.InetSocketAddress;
import java.time.Duration;

409
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java
View File

@@ -1,409 +0,0 @@
/*
* Copyright 2023-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.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;
import com.datastax.oss.driver.api.querybuilder.delete.Delete;
import com.datastax.oss.driver.api.querybuilder.delete.DeleteSelection;
import com.datastax.oss.driver.api.querybuilder.insert.InsertInto;
import com.datastax.oss.driver.api.querybuilder.insert.RegularInsert;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import io.micrometer.observation.ObservationRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentMetadata;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.model.EmbeddingUtils;
import org.springframework.ai.observation.conventions.VectorStoreProvider;
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext.Builder;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
/**
* The CassandraVectorStore is for managing and querying vector data in an Apache
* Cassandra db. It offers functionalities like adding, deleting, and performing
* similarity searches on documents.
*
* The store utilizes CQL to index and search vector data. It allows for custom metadata
* 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, column
* names, etc. It also requires an EmbeddingModel 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 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
* place, creating embeddings next to it, and performing vector similarity searches
* in-situ.
*
* 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
* embeddingModel 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
* @author Christian Tzolov
* @author Thomas Vitale
* @author Soby Chacko
* @see VectorStore
* @see org.springframework.ai.vectorstore.CassandraVectorStoreConfig
* @see EmbeddingModel
* @since 1.0.0
*/
public class CassandraVectorStore extends AbstractObservationVectorStore implements AutoCloseable {
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 static Map<Similarity, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(Similarity.COSINE,
VectorStoreSimilarityMetric.COSINE, Similarity.EUCLIDEAN, VectorStoreSimilarityMetric.EUCLIDEAN,
Similarity.DOT_PRODUCT, VectorStoreSimilarityMetric.DOT);
private final CassandraVectorStoreConfig conf;
private final EmbeddingModel embeddingModel;
private final FilterExpressionConverter filterExpressionConverter;
private final ConcurrentMap<Set<String>, PreparedStatement> addStmts = new ConcurrentHashMap<>();
private final PreparedStatement deleteStmt;
private final String similarityStmt;
private final Similarity similarity;
private final BatchingStrategy batchingStrategy;
public CassandraVectorStore(CassandraVectorStoreConfig conf, EmbeddingModel embeddingModel) {
this(conf, embeddingModel, ObservationRegistry.NOOP, null, new TokenCountBatchingStrategy());
}
public CassandraVectorStore(CassandraVectorStoreConfig conf, EmbeddingModel embeddingModel,
ObservationRegistry observationRegistry, VectorStoreObservationConvention customObservationConvention,
BatchingStrategy batchingStrategy) {
super(observationRegistry, customObservationConvention);
Preconditions.checkArgument(null != conf, "Config must not be null");
Preconditions.checkArgument(null != embeddingModel, "Embedding model must not be null");
this.conf = conf;
this.embeddingModel = embeddingModel;
conf.ensureSchemaExists(embeddingModel.dimensions());
prepareAddStatement(Set.of());
this.deleteStmt = prepareDeleteStatement();
TableMetadata cassandraMetadata = conf.session.getMetadata()
.getKeyspace(conf.schema.keyspace())
.get()
.getTable(conf.schema.table())
.get();
this.similarity = getIndexSimilarity(cassandraMetadata);
this.similarityStmt = similaritySearchStatement();
this.filterExpressionConverter = new CassandraFilterExpressionConverter(
cassandraMetadata.getColumns().values());
this.batchingStrategy = batchingStrategy;
}
private static Float[] toFloatArray(float[] embedding) {
Float[] embeddingFloat = new Float[embedding.length];
int i = 0;
for (Float d : embedding) {
embeddingFloat[i++] = d.floatValue();
}
return embeddingFloat;
}
@Override
public void doAdd(List<Document> documents) {
var futures = new CompletableFuture[documents.size()];
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
int i = 0;
for (Document d : documents) {
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());
}
builder = builder.setString(this.conf.schema.content(), d.getContent())
.setVector(this.conf.schema.embedding(),
CqlVector.newInstance(EmbeddingUtils.toList(embeddings.get(documents.indexOf(d)))),
Float.class);
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();
}
@Override
public Optional<Boolean> doDelete(List<String> idList) {
CompletableFuture[] futures = new CompletableFuture[idList.size()];
int i = 0;
for (String id : idList) {
List<Object> primaryKeyValues = this.conf.documentIdTranslator.apply(id);
BoundStatement s = this.deleteStmt.bind(primaryKeyValues.toArray());
futures[i++] = this.conf.session.executeAsync(s).toCompletableFuture();
}
CompletableFuture.allOf(futures).join();
return Optional.of(Boolean.TRUE);
}
@Override
public List<Document> doSimilaritySearch(SearchRequest request) {
Preconditions.checkArgument(request.getTopK() <= 1000);
var embedding = toFloatArray(this.embeddingModel.embed(request.getQuery()));
CqlVector<Float> cqlVector = CqlVector.newInstance(embedding);
String whereClause = "";
if (request.hasFilterExpression()) {
String expression = this.filterExpressionConverter.convertExpression(request.getFilterExpression());
if (!expression.isBlank()) {
whereClause = String.format("where %s", expression);
}
}
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(s)) {
float score = row.getFloat(0);
if (score < request.getSimilarityThreshold()) {
break;
}
Map<String, Object> docFields = new HashMap<>();
docFields.put(DocumentMetadata.DISTANCE.value(), 1 - score);
for (var metadata : this.conf.schema.metadataColumns()) {
var value = row.get(metadata.name(), metadata.javaType());
if (null != value) {
docFields.put(metadata.name(), value);
}
}
Document doc = Document.builder()
.id(getDocumentId(row))
.text(row.getString(this.conf.schema.content()))
.metadata(docFields)
.score((double) score)
.build();
documents.add(doc);
}
return documents;
}
@Override
public void close() throws Exception {
this.conf.close();
}
void checkSchemaValid() {
this.conf.checkSchemaValid(this.embeddingModel.dimensions());
}
private Similarity getIndexSimilarity(TableMetadata metadata) {
return Similarity.valueOf(metadata.getIndex(this.conf.schema.index())
.get()
.getOptions()
.getOrDefault("similarity_function", "COSINE")
.toUpperCase());
}
private PreparedStatement prepareDeleteStatement() {
Delete stmt = null;
DeleteSelection stmtStart = QueryBuilder.deleteFrom(this.conf.schema.keyspace(), this.conf.schema.table());
for (var c : this.conf.schema.partitionKeys()) {
stmt = (null != stmt ? stmt : stmtStart).whereColumn(c.name()).isEqualTo(QueryBuilder.bindMarker(c.name()));
}
for (var c : this.conf.schema.clusteringKeys()) {
stmt = stmt.whereColumn(c.name()).isEqualTo(QueryBuilder.bindMarker(c.name()));
}
return this.conf.session.prepare(stmt.build());
}
private PreparedStatement prepareAddStatement(Set<String> 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());
for (var c : this.conf.schema.partitionKeys()) {
stmt = (null != stmt ? stmt : stmtStart).value(c.name(), QueryBuilder.bindMarker(c.name()));
}
for (var c : this.conf.schema.clusteringKeys()) {
stmt = stmt.value(c.name(), QueryBuilder.bindMarker(c.name()));
}
stmt = stmt.value(this.conf.schema.content(), QueryBuilder.bindMarker(this.conf.schema.content()))
.value(this.conf.schema.embedding(), QueryBuilder.bindMarker(this.conf.schema.embedding()));
for (String metadataField : fields) {
stmt = stmt.value(metadataField, QueryBuilder.bindMarker(metadataField));
}
return this.conf.session.prepare(stmt.build());
});
}
private String similaritySearchStatement() {
StringBuilder ids = new StringBuilder();
for (var m : this.conf.schema.partitionKeys()) {
ids.append(m.name()).append(',');
}
for (var m : this.conf.schema.clusteringKeys()) {
ids.append(m.name()).append(',');
}
ids.deleteCharAt(ids.length() - 1);
String similarityFunction = new StringBuilder("similarity_").append(this.similarity.toString().toLowerCase())
.append('(')
.append(this.conf.schema.embedding())
.append(",?)")
.toString();
StringBuilder extraSelectFields = new StringBuilder();
for (var m : this.conf.schema.metadataColumns()) {
extraSelectFields.append(',').append(m.name());
}
if (this.conf.returnEmbeddings) {
extraSelectFields.append(',').append(this.conf.schema.embedding());
}
// java-driver-query-builder doesn't support orderByAnnOf yet
String query = String.format(QUERY_FORMAT, similarityFunction, ids.toString(), this.conf.schema.content(),
extraSelectFields.toString(), this.conf.schema.keyspace(), this.conf.schema.table(),
this.conf.schema.embedding());
query = query.replace("?", "%s");
logger.debug("preparing {}", query);
return query;
}
private String getDocumentId(Row row) {
List<Object> primaryKeyValues = new ArrayList<>();
for (var m : this.conf.schema.partitionKeys()) {
primaryKeyValues.add(row.get(m.name(), m.javaType()));
}
for (var m : this.conf.schema.clusteringKeys()) {
primaryKeyValues.add(row.get(m.name(), m.javaType()));
}
return this.conf.primaryKeyTranslator.apply(primaryKeyValues);
}
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.CASSANDRA.value(), operationName)
.withCollectionName(this.conf.schema.table())
.withDimensions(this.embeddingModel.dimensions())
.withNamespace(this.conf.schema.keyspace())
.withSimilarityMetric(getSimilarityMetric());
}
private String getSimilarityMetric() {
if (!SIMILARITY_TYPE_MAPPING.containsKey(this.similarity)) {
return this.similarity.name();
}
return SIMILARITY_TYPE_MAPPING.get(this.similarity).value();
}
/**
* Indexes are automatically created with COSINE. This can be changed manually via
* cqlsh
*/
public enum Similarity {
COSINE, DOT_PRODUCT, EUCLIDEAN
}
}

2
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java → vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/cassandra/CassandraFilterExpressionConverter.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
package org.springframework.ai.vectorstore.cassandra;
import java.util.Collection;
import java.util.Map;

1087
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStore.java Normal file
View File

File diff suppressed because it is too large Load Diff

29
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java → vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreConfig.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
package org.springframework.ai.vectorstore.cassandra;
import java.net.InetSocketAddress;
import java.util.Collections;
@@ -65,7 +65,9 @@ import org.springframework.lang.Nullable;
*
* @author Mick Semb Wever
* @since 1.0.0
* @deprecated since 1.0.0-M5, use {@link CassandraVectorStore#builder()} instead
*/
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public final class CassandraVectorStoreConfig implements AutoCloseable {
public static final String DEFAULT_KEYSPACE_NAME = "springframework";
@@ -317,6 +319,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
}
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public enum SchemaColumnTags {
INDEXED
@@ -329,20 +332,24 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
* It is a requirement that an empty {@code List<Object>} returns an example formatted
* id
*/
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public interface DocumentIdTranslator extends Function<String, List<Object>> {
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
/** Given a list of primary key column values, return the document id. */
public interface PrimaryKeyTranslator extends Function<List<Object>, String> {
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
record Schema(String keyspace, String table, List<SchemaColumn> partitionKeys, List<SchemaColumn> clusteringKeys,
String content, String embedding, String index, Set<SchemaColumn> metadataColumns) {
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public record SchemaColumn(String name, DataType type, SchemaColumnTags... tags) {
public SchemaColumn(String name, DataType type) {
@@ -364,6 +371,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public static final class Builder {
private CqlSession session = null;
@@ -405,6 +413,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
private Builder() {
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withCqlSession(CqlSession session) {
Preconditions.checkState(null == this.sessionBuilder,
"Cannot call withContactPoint(..) or withLocalDatacenter(..) and this method");
@@ -413,6 +422,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder addContactPoint(InetSocketAddress contactPoint) {
Preconditions.checkState(null == this.session, "Cannot call withCqlSession(..) and this method");
if (null == this.sessionBuilder) {
@@ -422,6 +432,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withLocalDatacenter(String localDC) {
Preconditions.checkState(null == this.session, "Cannot call withCqlSession(..) and this method");
if (null == this.sessionBuilder) {
@@ -431,22 +442,26 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withKeyspaceName(String keyspace) {
this.keyspace = keyspace;
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withTableName(String table) {
this.table = table;
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withPartitionKeys(List<SchemaColumn> partitionKeys) {
Preconditions.checkArgument(!partitionKeys.isEmpty());
this.partitionKeys = partitionKeys;
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withClusteringKeys(List<SchemaColumn> clusteringKeys) {
this.clusteringKeys = clusteringKeys;
return this;
@@ -456,21 +471,25 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
* defaults (if null) to '&lt;table_name&gt;_&lt;embedding_column_name&gt;_idx'
**/
@Nullable
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withIndexName(String name) {
this.indexName = name;
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withContentColumnName(String name) {
this.contentColumnName = name;
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withEmbeddingColumnName(String name) {
this.embeddingColumnName = name;
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder addMetadataColumns(SchemaColumn... columns) {
Builder builder = this;
for (SchemaColumn f : columns) {
@@ -479,12 +498,14 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
return builder;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder addMetadataColumns(List<SchemaColumn> columns) {
Builder builder = this;
this.metadataColumns.addAll(columns);
return builder;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder addMetadataColumn(SchemaColumn column) {
Preconditions.checkArgument(this.metadataColumns.stream().noneMatch(sc -> sc.name().equals(column.name())),
@@ -494,11 +515,13 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder disallowSchemaChanges() {
this.disallowSchemaChanges = true;
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder returnEmbeddings() {
this.returnEmbeddings = true;
return this;
@@ -510,22 +533,26 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
* utilize network. For local transformers you probably want a lower value to
* avoid saturation.
**/
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withFixedThreadPoolExecutorSize(int threads) {
Preconditions.checkArgument(0 < threads);
this.fixedThreadPoolExecutorSize = threads;
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withDocumentIdTranslator(DocumentIdTranslator documentIdTranslator) {
this.documentIdTranslator = documentIdTranslator;
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public Builder withPrimaryKeyTranslator(PrimaryKeyTranslator primaryKeyTranslator) {
this.primaryKeyTranslator = primaryKeyTranslator;
return this;
}
@Deprecated(since = "1.0.0-M5", forRemoval = true)
public CassandraVectorStoreConfig build() {
if (null == this.indexName) {
this.indexName = String.format("%s_%s_%s", this.table, this.embeddingColumnName, DEFAULT_INDEX_SUFFIX);

2
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/CassandraImage.java → vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/cassandra/CassandraImage.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai;
package org.springframework.ai.cassandra;
import org.testcontainers.utility.DockerImageName;

4
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/chat/memory/CassandraChatMemoryIT.java → vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/chat/memory/cassandra/CassandraChatMemoryIT.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.chat.memory;
package org.springframework.ai.chat.memory.cassandra;
import java.time.Duration;
@@ -26,7 +26,7 @@ import org.testcontainers.containers.CassandraContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.CassandraImage;
import org.springframework.ai.cassandra.CassandraImage;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;

2
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverterTests.java → vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraFilterExpressionConverterTests.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
package org.springframework.ai.vectorstore.cassandra;
import java.util.Collection;
import java.util.HashSet;

161
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java → vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraRichSchemaVectorStoreIT.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
package org.springframework.ai.vectorstore.cassandra;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
@@ -43,11 +43,12 @@ import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.shaded.org.apache.commons.lang3.RandomStringUtils;
import org.springframework.ai.CassandraImage;
import org.springframework.ai.cassandra.CassandraImage;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.transformers.TransformersEmbeddingModel;
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumn;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
@@ -92,66 +93,65 @@ class CassandraRichSchemaVectorStoreIT {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class);
static CassandraVectorStoreConfig.Builder storeBuilder(ApplicationContext context,
List<SchemaColumn> columnOverrides) throws IOException {
static CassandraVectorStore.CassandraBuilder storeBuilder(ApplicationContext context,
List<CassandraVectorStore.SchemaColumn> columnOverrides) throws IOException {
Optional<SchemaColumn> wikiOverride = columnOverrides.stream().filter(f -> "wiki".equals(f.name())).findFirst();
Optional<CassandraVectorStore.SchemaColumn> wikiOverride = columnOverrides.stream()
.filter(f -> "wiki".equals(f.name()))
.findFirst();
Optional<SchemaColumn> langOverride = columnOverrides.stream()
Optional<CassandraVectorStore.SchemaColumn> langOverride = columnOverrides.stream()
.filter(f -> "language".equals(f.name()))
.findFirst();
Optional<SchemaColumn> titleOverride = columnOverrides.stream()
Optional<CassandraVectorStore.SchemaColumn> titleOverride = columnOverrides.stream()
.filter(f -> "title".equals(f.name()))
.findFirst();
Optional<SchemaColumn> chunkNoOverride = columnOverrides.stream()
Optional<CassandraVectorStore.SchemaColumn> chunkNoOverride = columnOverrides.stream()
.filter(f -> "chunk_no".equals(f.name()))
.findFirst();
SchemaColumn wikiSC = wikiOverride.orElse(new SchemaColumn("wiki", DataTypes.TEXT));
SchemaColumn langSC = langOverride.orElse(new SchemaColumn("language", DataTypes.TEXT));
SchemaColumn titleSC = titleOverride.orElse(new SchemaColumn("title", DataTypes.TEXT));
SchemaColumn chunkNoSC = chunkNoOverride.orElse(new SchemaColumn("chunk_no", DataTypes.INT));
var wikiSC = wikiOverride.orElse(new CassandraVectorStore.SchemaColumn("wiki", DataTypes.TEXT));
var langSC = langOverride.orElse(new CassandraVectorStore.SchemaColumn("language", DataTypes.TEXT));
var titleSC = titleOverride.orElse(new CassandraVectorStore.SchemaColumn("title", DataTypes.TEXT));
var chunkNoSC = chunkNoOverride.orElse(new CassandraVectorStore.SchemaColumn("chunk_no", DataTypes.INT));
List<SchemaColumn> partitionKeys = List.of(wikiSC, langSC, titleSC);
List<SchemaColumn> clusteringKeys = List.of(chunkNoSC);
CassandraVectorStoreConfig.Builder builder = CassandraVectorStoreConfig.builder()
.withCqlSession(context.getBean(CqlSession.class))
.withKeyspaceName("test_wikidata")
.withTableName("articles")
.withPartitionKeys(partitionKeys)
.withClusteringKeys(clusteringKeys)
.withContentColumnName("body")
.withEmbeddingColumnName("all_minilm_l6_v2_embedding")
.withIndexName("all_minilm_l6_v2_ann")
.addMetadataColumns(new SchemaColumn("revision", DataTypes.INT),
new SchemaColumn("id", DataTypes.INT, CassandraVectorStoreConfig.SchemaColumnTags.INDEXED))
List<CassandraVectorStore.SchemaColumn> partitionKeys = List.of(wikiSC, langSC, titleSC);
List<CassandraVectorStore.SchemaColumn> clusteringKeys = List.of(chunkNoSC);
return CassandraVectorStore.builder()
.session(context.getBean(CqlSession.class))
.keyspace("test_wikidata")
.table("articles")
.partitionKeys(partitionKeys)
.clusteringKeys(clusteringKeys)
.contentColumnName("body")
.embeddingColumnName("all_minilm_l6_v2_embedding")
.indexName("all_minilm_l6_v2_ann")
.addMetadataColumns(new CassandraVectorStore.SchemaColumn("revision", DataTypes.INT),
new CassandraVectorStore.SchemaColumn("id", DataTypes.INT,
CassandraVectorStore.SchemaColumnTags.INDEXED))
// this store uses '§' as a deliminator in the document id between db columns
// 'title' and 'chunk_no'
.withPrimaryKeyTranslator((List<Object> primaryKeys) -> {
.primaryKeyTranslator((List<Object> primaryKeys) -> {
if (primaryKeys.isEmpty()) {
return "test§¶0";
}
return java.lang.String.format("%s§¶%s", primaryKeys.get(2), primaryKeys.get(3));
return String.format("%s§¶%s", primaryKeys.get(2), primaryKeys.get(3));
})
.withDocumentIdTranslator(id -> {
.documentIdTranslator(id -> {
String[] parts = id.split("§¶");
String title = parts[0];
int chunk_no = 0 < parts.length ? Integer.parseInt(parts[1]) : 0;
return List.of("simplewiki", "en", title, chunk_no);
});
return builder;
}
@Test
void ensureSchemaCreation() {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
try (CassandraVectorStore store = createStore(context, false)) {
Assertions.assertNotNull(store);
store.checkSchemaValid();
store.similaritySearch(SearchRequest.query("1843").withTopK(1));
@@ -163,14 +163,16 @@ class CassandraRichSchemaVectorStoreIT {
void ensureSchemaNoCreation() {
this.contextRunner.run(context -> {
executeCqlFile(context, "test_wiki_full_schema.cql");
var wrapper = createStore(context, List.of(), true, false);
var builder = createBuilder(context, List.of(), true, false);
Assertions.assertNotNull(builder);
var store = new CassandraVectorStore(builder);
try {
Assertions.assertNotNull(wrapper.store());
wrapper.store().checkSchemaValid();
wrapper.store().similaritySearch(SearchRequest.query("1843").withTopK(1));
store.checkSchemaValid();
wrapper.conf().dropKeyspace();
store.similaritySearch(SearchRequest.query("1843").withTopK(1));
CassandraVectorStore.dropKeyspace(builder);
executeCqlFile(context, "test_wiki_partial_3_schema.cql");
// IllegalStateException: column all_minilm_l6_v2_embedding does not exist
@@ -180,8 +182,8 @@ class CassandraRichSchemaVectorStoreIT {
Assertions.assertEquals("column all_minilm_l6_v2_embedding does not exist", ise.getMessage());
}
finally {
wrapper.conf().dropKeyspace();
wrapper.store().close();
CassandraVectorStore.dropKeyspace(builder);
store.close();
}
});
}
@@ -192,16 +194,18 @@ class CassandraRichSchemaVectorStoreIT {
int PARTIAL_FILES = 5;
for (int i = 0; i < PARTIAL_FILES; ++i) {
executeCqlFile(context, java.lang.String.format("test_wiki_partial_%d_schema.cql", i));
var wrapper = createStore(context, List.of(), false, false);
var builder = createBuilder(context, List.of(), false, false);
Assertions.assertNotNull(builder);
CassandraVectorStore.dropKeyspace(builder);
var store = builder.build();
try {
Assertions.assertNotNull(wrapper.store());
wrapper.store().checkSchemaValid();
store.checkSchemaValid();
wrapper.store().similaritySearch(SearchRequest.query("1843").withTopK(1));
wrapper.conf().dropKeyspace();
store.similaritySearch(SearchRequest.query("1843").withTopK(1));
}
finally {
wrapper.store().close();
CassandraVectorStore.dropKeyspace(builder);
store.close();
}
}
// make sure there's not more files to test
@@ -213,7 +217,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void addAndSearch() {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
try (CassandraVectorStore store = createStore(context, false)) {
store.add(documents);
List<Document> results = store
@@ -241,16 +245,16 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void addAndSearchPoormansBench() {
// todo – replace with JMH (parameters: nThreads, rounds, runs, docsPerAdd)
int nThreads = CassandraVectorStoreConfig.DEFAULT_ADD_CONCURRENCY;
int nThreads = CassandraVectorStore.DEFAULT_ADD_CONCURRENCY;
int runs = 10; // 100;
int docsPerAdd = 12; // 128;
int rounds = 3;
this.contextRunner.run(context -> {
try (CassandraVectorStore store = new CassandraVectorStore(
storeBuilder(context, List.of()).withFixedThreadPoolExecutorSize(nThreads).build(),
context.getBean(EmbeddingModel.class))) {
try (CassandraVectorStore store = storeBuilder(context, List.of()).fixedThreadPoolExecutorSize(nThreads)
.embeddingModel(context.getBean(EmbeddingModel.class))
.build()) {
var executor = Executors.newFixedThreadPool((int) (nThreads * 1.2));
for (int k = 0; k < rounds; ++k) {
@@ -286,7 +290,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void searchWithPartitionFilter() throws InterruptedException {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
try (CassandraVectorStore store = createStore(context, false)) {
store.add(documents);
List<Document> results = store.similaritySearch(SearchRequest.query("Great Dark Spot").withTopK(5));
@@ -336,7 +340,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void unsearchableFilters() throws InterruptedException {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
try (CassandraVectorStore store = createStore(context, false)) {
store.add(documents);
List<Document> results = store.similaritySearch(SearchRequest.query("Great Dark Spot").withTopK(5));
@@ -354,7 +358,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void searchWithFilters() throws InterruptedException {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
try (CassandraVectorStore store = createStore(context, false)) {
store.add(documents);
List<Document> results = store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY).withTopK(5));
@@ -418,10 +422,10 @@ class CassandraRichSchemaVectorStoreIT {
this.contextRunner.run(context -> {
List<SchemaColumn> overrides = List.of(
new SchemaColumn("title", DataTypes.TEXT, CassandraVectorStoreConfig.SchemaColumnTags.INDEXED),
new SchemaColumn("chunk_no", DataTypes.INT, CassandraVectorStoreConfig.SchemaColumnTags.INDEXED));
new SchemaColumn("title", DataTypes.TEXT, CassandraVectorStore.SchemaColumnTags.INDEXED),
new SchemaColumn("chunk_no", DataTypes.INT, CassandraVectorStore.SchemaColumnTags.INDEXED));
try (CassandraVectorStore store = createStore(context, overrides, false, true).store()) {
try (CassandraVectorStore store = createStore(context, overrides, false, true)) {
store.add(documents);
@@ -452,7 +456,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void documentUpdate() {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
try (CassandraVectorStore store = createStore(context, false)) {
store.add(documents);
List<Document> results = store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY).withTopK(1));
@@ -502,7 +506,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void searchWithThreshold() {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
try (CassandraVectorStore store = createStore(context, false)) {
store.add(documents);
List<Document> fullResult = store
@@ -530,26 +534,43 @@ class CassandraRichSchemaVectorStoreIT {
});
}
private StoreWrapper<CassandraVectorStore, CassandraVectorStoreConfig> createStore(ApplicationContext context,
boolean disallowSchemaCreation) throws IOException {
private CassandraVectorStore createStore(ApplicationContext context, boolean disallowSchemaCreation)
throws IOException {
return createStore(context, List.of(), disallowSchemaCreation, true);
}
private StoreWrapper<CassandraVectorStore, CassandraVectorStoreConfig> createStore(ApplicationContext context,
private CassandraVectorStore createStore(ApplicationContext context, List<SchemaColumn> columnOverrides,
boolean disallowSchemaCreation, boolean dropKeyspaceFirst) throws IOException {
CassandraVectorStore.CassandraBuilder builder = storeBuilder(context, columnOverrides);
if (disallowSchemaCreation) {
builder = builder.disallowSchemaChanges(true);
}
if (dropKeyspaceFirst) {
CassandraVectorStore.dropKeyspace(builder);
}
builder.embeddingModel(context.getBean(EmbeddingModel.class));
return new CassandraVectorStore(builder);
}
private CassandraVectorStore.CassandraBuilder createBuilder(ApplicationContext context,
List<SchemaColumn> columnOverrides, boolean disallowSchemaCreation, boolean dropKeyspaceFirst)
throws IOException {
CassandraVectorStoreConfig.Builder builder = storeBuilder(context, columnOverrides);
CassandraVectorStore.CassandraBuilder builder = storeBuilder(context, columnOverrides);
if (disallowSchemaCreation) {
builder = builder.disallowSchemaChanges();
builder = builder.disallowSchemaChanges(true);
}
CassandraVectorStoreConfig conf = builder.build();
if (dropKeyspaceFirst) {
conf.dropKeyspace();
CassandraVectorStore.dropKeyspace(builder);
}
return new StoreWrapper(new CassandraVectorStore(conf, context.getBean(EmbeddingModel.class)), conf);
builder.embeddingModel(context.getBean(EmbeddingModel.class));
return builder;
}
private void executeCqlFile(ApplicationContext context, String filename) throws IOException {
@@ -588,8 +609,4 @@ class CassandraRichSchemaVectorStoreIT {
}
public record StoreWrapper<K, V>(K store, V conf) {
}
}

56
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java → vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreIT.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
package org.springframework.ai.vectorstore.cassandra;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
@@ -35,12 +35,13 @@ import org.testcontainers.containers.CassandraContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.CassandraImage;
import org.springframework.ai.cassandra.CassandraImage;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.transformers.TransformersEmbeddingModel;
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn;
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumnTags;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumn;
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumnTags;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
@@ -84,24 +85,25 @@ class CassandraVectorStoreIT {
}
}
private static CassandraVectorStoreConfig.Builder storeBuilder(CqlSession cqlSession) {
return CassandraVectorStoreConfig.builder()
.withCqlSession(cqlSession)
.withKeyspaceName("test_" + CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME);
private static CassandraVectorStore.CassandraBuilder storeBuilder(CqlSession cqlSession) {
return CassandraVectorStore.builder()
.session(cqlSession)
.keyspace("test_" + CassandraVectorStore.DEFAULT_KEYSPACE_NAME);
}
private static CassandraVectorStore createTestStore(ApplicationContext context, SchemaColumn... metadataFields) {
CassandraVectorStoreConfig.Builder builder = storeBuilder(context.getBean(CqlSession.class))
CassandraVectorStore.CassandraBuilder builder = storeBuilder(context.getBean(CqlSession.class))
.addMetadataColumns(metadataFields);
return createTestStore(context, builder);
}
private static CassandraVectorStore createTestStore(ApplicationContext context,
CassandraVectorStoreConfig.Builder builder) {
CassandraVectorStoreConfig conf = builder.build();
conf.dropKeyspace();
return new CassandraVectorStore(conf, context.getBean(EmbeddingModel.class));
CassandraVectorStore.CassandraBuilder builder) {
CassandraVectorStore.dropKeyspace(builder);
builder.embeddingModel(context.getBean(EmbeddingModel.class));
CassandraVectorStore store = builder.build();
return store;
}
@Test
@@ -147,8 +149,8 @@ class CassandraVectorStoreIT {
@Test
void addAndSearchReturnEmbeddings() {
this.contextRunner.run(context -> {
CassandraVectorStoreConfig.Builder builder = storeBuilder(context.getBean(CqlSession.class))
.returnEmbeddings();
CassandraVectorStore.CassandraBuilder builder = storeBuilder(context.getBean(CqlSession.class))
.returnEmbeddings(true);
try (CassandraVectorStore store = createTestStore(context, builder)) {
List<Document> documents = documents();
@@ -197,8 +199,7 @@ class CassandraVectorStoreIT {
results = store.similaritySearch(SearchRequest.query("The World")
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression(
java.lang.String.format("%s == 'NL'", CassandraVectorStoreConfig.DEFAULT_ID_NAME)));
.withFilterExpression(java.lang.String.format("%s == 'NL'", CassandraVectorStore.DEFAULT_ID_NAME)));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
@@ -207,7 +208,7 @@ class CassandraVectorStoreIT {
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression(
java.lang.String.format("%s == 'BG2'", CassandraVectorStoreConfig.DEFAULT_ID_NAME)));
java.lang.String.format("%s == 'BG2'", CassandraVectorStore.DEFAULT_ID_NAME)));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument2.getId());
@@ -216,7 +217,7 @@ class CassandraVectorStoreIT {
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression(java.lang.String.format("%s == 'BG' && year == 2020",
CassandraVectorStoreConfig.DEFAULT_ID_NAME)));
CassandraVectorStore.DEFAULT_ID_NAME)));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
@@ -227,7 +228,7 @@ class CassandraVectorStoreIT {
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression(java.lang.String.format("NOT(%s == 'BG' && year == 2020)",
CassandraVectorStoreConfig.DEFAULT_ID_NAME))));
CassandraVectorStore.DEFAULT_ID_NAME))));
}
});
}
@@ -394,14 +395,15 @@ class CassandraVectorStoreIT {
@Bean
public CassandraVectorStore store(CqlSession cqlSession, EmbeddingModel embeddingModel) {
CassandraVectorStoreConfig conf = storeBuilder(cqlSession)
.addMetadataColumns(new SchemaColumn("meta1", DataTypes.TEXT),
new SchemaColumn("meta2", DataTypes.TEXT), new SchemaColumn("country", DataTypes.TEXT),
new SchemaColumn("year", DataTypes.SMALLINT))
.build();
CassandraVectorStore.CassandraBuilder builder = storeBuilder(cqlSession)
.addMetadataColumns(new CassandraVectorStore.SchemaColumn("meta1", DataTypes.TEXT),
new CassandraVectorStore.SchemaColumn("meta2", DataTypes.TEXT),
new CassandraVectorStore.SchemaColumn("country", DataTypes.TEXT),
new CassandraVectorStore.SchemaColumn("year", DataTypes.SMALLINT))
.embeddingModel(embeddingModel);
conf.dropKeyspace();
return new CassandraVectorStore(conf, embeddingModel);
CassandraVectorStore.dropKeyspace(builder);
return builder.build();
}
@Bean

38
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreObservationIT.java → vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreObservationIT.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
package org.springframework.ai.vectorstore.cassandra;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
@@ -32,7 +32,7 @@ import org.testcontainers.containers.CassandraContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.CassandraImage;
import org.springframework.ai.cassandra.CassandraImage;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
@@ -40,7 +40,8 @@ import org.springframework.ai.observation.conventions.SpringAiKind;
import org.springframework.ai.observation.conventions.VectorStoreProvider;
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
import org.springframework.ai.transformers.TransformersEmbeddingModel;
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames;
@@ -80,12 +81,6 @@ public class CassandraVectorStoreObservationIT {
}
}
private static CassandraVectorStoreConfig.Builder storeBuilder(CqlSession cqlSession) {
return CassandraVectorStoreConfig.builder()
.withCqlSession(cqlSession)
.withKeyspaceName("test_" + CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME);
}
@Test
void observationVectorStoreAddAndQueryOperations() {
@@ -110,7 +105,7 @@ public class CassandraVectorStoreObservationIT {
.doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_VECTOR_QUERY_CONTENT.asString())
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_DIMENSION_COUNT.asString(), "384")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_COLLECTION_NAME.asString(),
CassandraVectorStoreConfig.DEFAULT_TABLE_NAME)
CassandraVectorStore.DEFAULT_TABLE_NAME)
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_NAMESPACE.asString(), "test_springframework")
.doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_VECTOR_FIELD_NAME.asString())
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_SEARCH_SIMILARITY_METRIC.asString(),
@@ -144,7 +139,7 @@ public class CassandraVectorStoreObservationIT {
"What is Great Depression")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_DIMENSION_COUNT.asString(), "384")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_COLLECTION_NAME.asString(),
CassandraVectorStoreConfig.DEFAULT_TABLE_NAME)
CassandraVectorStore.DEFAULT_TABLE_NAME)
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_NAMESPACE.asString(), "test_springframework")
.doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_VECTOR_FIELD_NAME.asString())
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_SEARCH_SIMILARITY_METRIC.asString(),
@@ -177,15 +172,20 @@ public class CassandraVectorStoreObservationIT {
public CassandraVectorStore store(CqlSession cqlSession, EmbeddingModel embeddingModel,
ObservationRegistry observationRegistry) {
CassandraVectorStoreConfig conf = storeBuilder(cqlSession)
.addMetadataColumns(new SchemaColumn("meta1", DataTypes.TEXT),
new SchemaColumn("meta2", DataTypes.TEXT), new SchemaColumn("country", DataTypes.TEXT),
new SchemaColumn("year", DataTypes.SMALLINT))
.build();
CassandraVectorStore.CassandraBuilder builder = CassandraVectorStore.builder()
.session(cqlSession)
.session(cqlSession)
.keyspace("test_" + CassandraVectorStore.DEFAULT_KEYSPACE_NAME)
.addMetadataColumns(new CassandraVectorStore.SchemaColumn("meta1", DataTypes.TEXT),
new CassandraVectorStore.SchemaColumn("meta2", DataTypes.TEXT),
new CassandraVectorStore.SchemaColumn("country", DataTypes.TEXT),
new CassandraVectorStore.SchemaColumn("year", DataTypes.SMALLINT))
.embeddingModel(embeddingModel)
.observationRegistry(observationRegistry)
.batchingStrategy(new TokenCountBatchingStrategy());
conf.dropKeyspace();
return new CassandraVectorStore(conf, embeddingModel, observationRegistry, null,
new TokenCountBatchingStrategy());
CassandraVectorStore.dropKeyspace(builder);
return builder.build();
}
@Bean

38
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java → vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/WikiVectorStoreExample.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
package org.springframework.ai.vectorstore.cassandra;
import java.util.List;
@@ -22,12 +22,14 @@ import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
import com.datastax.oss.driver.api.core.type.DataTypes;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.transformers.TransformersEmbeddingModel;
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumn;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
@@ -46,6 +48,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* @since 1.0.0
*/
@Testcontainers
@Disabled("This is an example, not a really a test as it requires external setup")
class WikiVectorStoreExample {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
@@ -90,36 +93,33 @@ class WikiVectorStoreExample {
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(partitionColumns)
.withClusteringKeys(clusteringColumns)
.withContentColumnName("body")
.withEmbeddingColumnName("all_minilm_l6_v2_embedding")
.withIndexName("all_minilm_l6_v2_ann")
.disallowSchemaChanges()
return CassandraVectorStore.builder()
.session(cqlSession)
.keyspace("wikidata")
.table("articles")
.partitionKeys(partitionColumns)
.clusteringKeys(clusteringColumns)
.contentColumnName("body")
.embeddingColumnName("all_minilm_l6_v2_embedding")
.indexName("all_minilm_l6_v2_ann")
.disallowSchemaChanges(true)
.addMetadataColumns(extraColumns)
.withPrimaryKeyTranslator((List<Object> primaryKeys) -> {
.primaryKeyTranslator((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";
}
return java.lang.String.format("%s§¶%s", primaryKeys.get(2), primaryKeys.get(3));
return String.format("%s§¶%s", primaryKeys.get(2), primaryKeys.get(3));
})
.withDocumentIdTranslator(id -> {
.documentIdTranslator(id -> {
String[] parts = id.split("§¶");
String title = parts[0];
int chunk_no = 0 < parts.length ? Integer.parseInt(parts[1]) : 0;
return List.of("simplewiki", "en", title, chunk_no, 0);
})
.embeddingModel(embeddingModel())
.build();
return new CassandraVectorStore(conf, embeddingModel());
}
@Bean