Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -74,8 +74,7 @@ public CassandraVectorStore vectorStore(EmbeddingModel embeddingModel, Cassandra
.embeddingColumnName(properties.getEmbeddingColumnName())
.indexName(properties.getIndexName())
.fixedThreadPoolExecutorSize(properties.getFixedThreadPoolExecutorSize())
.disallowSchemaChanges(!properties.isInitializeSchema())
.returnEmbeddings(properties.getReturnEmbeddings())
.initializeSchema(properties.isInitializeSchema())
.observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
.customObservationConvention(customObservationConvention.getIfAvailable(() -> null))
.batchingStrategy(batchingStrategy)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,6 @@

package org.springframework.ai.vectorstore.cassandra.autoconfigure;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore;
import org.springframework.ai.vectorstore.properties.CommonVectorStoreProperties;
import org.springframework.boot.context.properties.ConfigurationProperties;
Expand All @@ -35,8 +32,6 @@ public class CassandraVectorStoreProperties extends CommonVectorStoreProperties

public static final String CONFIG_PREFIX = "spring.ai.vectorstore.cassandra";

private static final Logger logger = LoggerFactory.getLogger(CassandraVectorStoreProperties.class);

private String keyspace = CassandraVectorStore.DEFAULT_KEYSPACE_NAME;

private String table = CassandraVectorStore.DEFAULT_TABLE_NAME;
Expand All @@ -47,8 +42,6 @@ public class CassandraVectorStoreProperties extends CommonVectorStoreProperties

private String embeddingColumnName = CassandraVectorStore.DEFAULT_EMBEDDING_COLUMN_NAME;

private boolean returnEmbeddings = false;

private int fixedThreadPoolExecutorSize = CassandraVectorStore.DEFAULT_ADD_CONCURRENCY;

public String getKeyspace() {
Expand Down Expand Up @@ -91,14 +84,6 @@ public void setEmbeddingColumnName(String embeddingColumnName) {
this.embeddingColumnName = embeddingColumnName;
}

public boolean getReturnEmbeddings() {
return this.returnEmbeddings;
}

public void setReturnEmbeddings(boolean returnEmbeddings) {
this.returnEmbeddings = returnEmbeddings;
}

public int getFixedThreadPoolExecutorSize() {
return this.fixedThreadPoolExecutorSize;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,9 @@ This Spring AI Vector Store is designed to work for both brand-new RAG applicati

The store can also be used for non-RAG use-cases in an existing database, e.g. semantic searches, geo-proximity searches, etc.

The 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`.
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 `initializeSchema`.

When using spring-boot-autoconfigure `disallowSchemaChanges` defaults to true, per Spring Boot standards, and you must opt-in to schema creation/modifications by setting `...initialize-schema=true` in the `application.properties` file.
When using spring-boot-autoconfigure `initializeSchema` defaults to `false`, per Spring Boot standards, and you must opt-in to schema creation/modifications by setting `...initialize-schema=true` in the `application.properties` file.

== What is JVector?

Expand Down Expand Up @@ -167,7 +167,7 @@ public VectorStore vectorStore(CqlSession session, EmbeddingModel embeddingModel
// Performance tuning
.fixedThreadPoolExecutorSize(32)
// Schema management
.disallowSchemaChanges(false)
.initializeSchema(true)
// Custom batching strategy
.batchingStrategy(new TokenCountBatchingStrategy())
.build();
Expand Down Expand Up @@ -282,7 +282,7 @@ public VectorStore vectorStore(CqlSession session, EmbeddingModel embeddingModel
.contentColumnName("body")
.embeddingColumnName("all_minilm_l6_v2_embedding")
.indexName("all_minilm_l6_v2_ann")
.disallowSchemaChanges(true)
.initializeSchema(false)
.addMetadataColumns(extraColumns)
.primaryKeyTranslator((List<Object> primaryKeys) -> {
if (primaryKeys.isEmpty()) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@
*
* 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 CassandraBuilder#disallowSchemaChanges().
* Disable this with the CassandraBuilder#initializeSchema(boolean) method().
*
* <p>
* Basic usage example:
Expand Down Expand Up @@ -139,7 +139,7 @@
* .contentColumnName("text")
* .embeddingColumnName("vector")
* .fixedThreadPoolExecutorSize(32)
* .disallowSchemaChanges(false)
* .initializeSchema(true)
* .batchingStrategy(new TokenCountBatchingStrategy())
* .build();
* }</pre>
Expand Down Expand Up @@ -202,7 +202,7 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme

private final Schema schema;

private final boolean disallowSchemaChanges;
private final boolean initializeSchema;

private final FilterExpressionConverter filterExpressionConverter;

Expand All @@ -229,7 +229,7 @@ protected CassandraVectorStore(Builder builder) {

this.session = builder.session;
this.schema = builder.buildSchema();
this.disallowSchemaChanges = builder.disallowSchemaChanges;
this.initializeSchema = builder.initializeSchema;
this.documentIdTranslator = builder.documentIdTranslator;
this.primaryKeyTranslator = builder.primaryKeyTranslator;
this.executor = Executors.newFixedThreadPool(builder.fixedThreadPoolExecutorSize);
Expand Down Expand Up @@ -525,7 +525,7 @@ static void dropKeyspace(Builder builder) {
}

void ensureSchemaExists(int vectorDimension) {
if (!this.disallowSchemaChanges) {
if (this.initializeSchema) {
SchemaUtil.ensureKeyspaceExists(this.session, this.schema.keyspace);
ensureTableExists(vectorDimension);
ensureTableColumnsExist(vectorDimension);
Expand Down Expand Up @@ -805,7 +805,7 @@ public static class Builder extends AbstractVectorStoreBuilder<Builder> {

private Set<SchemaColumn> metadataColumns = new HashSet<>();

private boolean disallowSchemaChanges = false;
private boolean initializeSchema = true;

private int fixedThreadPoolExecutorSize = DEFAULT_ADD_CONCURRENCY;

Expand All @@ -821,8 +821,6 @@ public static class Builder extends AbstractVectorStoreBuilder<Builder> {
return (String) primaryKeyColumns.get(0);
};

private boolean returnEmbeddings = false;

private Builder(EmbeddingModel embeddingModel) {
super(embeddingModel);
}
Expand Down Expand Up @@ -938,12 +936,12 @@ public Builder indexName(String indexName) {
}

/**
* Sets whether to disallow schema changes.
* @param disallowSchemaChanges true to disallow schema changes
* Sets whether to initialize the schema.
* @param initializeSchema true to initialize schema, false otherwise
* @return the builder instance
*/
public Builder disallowSchemaChanges(boolean disallowSchemaChanges) {
this.disallowSchemaChanges = disallowSchemaChanges;
public Builder initializeSchema(boolean initializeSchema) {
this.initializeSchema = initializeSchema;
return this;
}

Expand Down Expand Up @@ -1016,11 +1014,6 @@ public Builder primaryKeyTranslator(PrimaryKeyTranslator translator) {
return this;
}

public Builder returnEmbeddings(boolean returnEmbeddings) {
this.returnEmbeddings = true;
return this;
}

Schema buildSchema() {
if (this.indexName == null) {
this.indexName = String.format("%s_%s_%s", this.table, this.embeddingColumnName, DEFAULT_INDEX_SUFFIX);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,7 @@ static CassandraVectorStore.Builder storeBuilder(ApplicationContext context,
@Test
void ensureSchemaCreation() {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false)) {
try (CassandraVectorStore store = createStore(context, true)) {
Assertions.assertNotNull(store);
store.checkSchemaValid();
store.similaritySearch(SearchRequest.builder().query("1843").topK(1).build());
Expand All @@ -162,7 +162,7 @@ void ensureSchemaCreation() {
void ensureSchemaNoCreation() {
this.contextRunner.run(context -> {
executeCqlFile(context, "test_wiki_full_schema.cql");
var builder = createBuilder(context, List.of(), true, false);
var builder = createBuilder(context, List.of(), false, false);
Assertions.assertNotNull(builder);
var store = new CassandraVectorStore(builder);
try {
Expand All @@ -176,7 +176,7 @@ void ensureSchemaNoCreation() {

// IllegalStateException: column all_minilm_l6_v2_embedding does not exist
IllegalStateException ise = Assertions.assertThrows(IllegalStateException.class,
() -> createStore(context, List.of(), true, false));
() -> createStore(context, List.of(), false, false));

Assertions.assertEquals("column all_minilm_l6_v2_embedding does not exist", ise.getMessage());
}
Expand All @@ -193,7 +193,7 @@ void ensureSchemaPartialCreation() {
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 builder = createBuilder(context, List.of(), false, false);
var builder = createBuilder(context, List.of(), true, false);
Assertions.assertNotNull(builder);
CassandraVectorStore.dropKeyspace(builder);
var store = builder.build();
Expand All @@ -216,7 +216,7 @@ void ensureSchemaPartialCreation() {
@Test
void addAndSearch() {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false)) {
try (CassandraVectorStore store = createStore(context, true)) {
store.add(documents);

List<Document> results = store.similaritySearch(
Expand Down Expand Up @@ -290,7 +290,7 @@ void addAndSearchPoormansBench() {
@Test
void searchWithPartitionFilter() throws InterruptedException {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false)) {
try (CassandraVectorStore store = createStore(context, true)) {
store.add(documents);

List<Document> results = store
Expand Down Expand Up @@ -346,7 +346,7 @@ void searchWithPartitionFilter() throws InterruptedException {
@Test
void unsearchableFilters() throws InterruptedException {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false)) {
try (CassandraVectorStore store = createStore(context, true)) {
store.add(documents);

List<Document> results = store
Expand All @@ -367,7 +367,7 @@ void unsearchableFilters() throws InterruptedException {
@Test
void searchWithFilters() throws InterruptedException {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false)) {
try (CassandraVectorStore store = createStore(context, true)) {
store.add(documents);

List<Document> results = store
Expand Down Expand Up @@ -447,7 +447,7 @@ void searchWithFilterOnPrimaryKeys() throws InterruptedException {
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)) {
try (CassandraVectorStore store = createStore(context, overrides, true, true)) {

store.add(documents);

Expand Down Expand Up @@ -481,7 +481,7 @@ void searchWithFilterOnPrimaryKeys() throws InterruptedException {
@Test
void documentUpdate() {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false)) {
try (CassandraVectorStore store = createStore(context, true)) {
store.add(documents);

List<Document> results = store
Expand Down Expand Up @@ -532,7 +532,7 @@ void documentUpdate() {
@Test
void searchWithThreshold() {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false)) {
try (CassandraVectorStore store = createStore(context, true)) {
store.add(documents);

List<Document> fullResult = store.similaritySearch(
Expand Down Expand Up @@ -562,19 +562,16 @@ void searchWithThreshold() {
});
}

private CassandraVectorStore createStore(ApplicationContext context, boolean disallowSchemaCreation)
throws IOException {
private CassandraVectorStore createStore(ApplicationContext context, boolean initializeSchema) throws IOException {

return createStore(context, List.of(), disallowSchemaCreation, true);
return createStore(context, List.of(), initializeSchema, true);
}

private CassandraVectorStore createStore(ApplicationContext context, List<SchemaColumn> columnOverrides,
boolean disallowSchemaCreation, boolean dropKeyspaceFirst) throws IOException {
boolean initializeSchema, boolean dropKeyspaceFirst) throws IOException {

CassandraVectorStore.Builder builder = storeBuilder(context, columnOverrides);
if (disallowSchemaCreation) {
builder = builder.disallowSchemaChanges(true);
}
builder.initializeSchema(initializeSchema);

if (dropKeyspaceFirst) {
CassandraVectorStore.dropKeyspace(builder);
Expand All @@ -584,12 +581,10 @@ private CassandraVectorStore createStore(ApplicationContext context, List<Schema
}

private CassandraVectorStore.Builder createBuilder(ApplicationContext context, List<SchemaColumn> columnOverrides,
boolean disallowSchemaCreation, boolean dropKeyspaceFirst) throws IOException {
boolean initailzeSchema, boolean dropKeyspaceFirst) throws IOException {

CassandraVectorStore.Builder builder = storeBuilder(context, columnOverrides);
if (disallowSchemaCreation) {
builder = builder.disallowSchemaChanges(true);
}
builder.initializeSchema(initailzeSchema);

if (dropKeyspaceFirst) {
CassandraVectorStore.dropKeyspace(builder);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -172,39 +172,6 @@ void addAndSearch() {
});
}

@Test
void addAndSearchReturnEmbeddings() {
this.contextRunner.run(context -> {
CassandraVectorStore.Builder builder = storeBuilder(context.getBean(CqlSession.class),
context.getBean(EmbeddingModel.class))
.returnEmbeddings(true);

try (CassandraVectorStore store = createTestStore(context, builder)) {
List<Document> documents = documents();
store.add(documents);

List<Document> results = store
.similaritySearch(SearchRequest.builder().query("Spring").topK(1).build());

assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents().get(0).getId());

assertThat(resultDoc.getText()).contains(
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");

assertThat(resultDoc.getMetadata()).hasSize(1);
assertThat(resultDoc.getMetadata()).containsKey(DocumentMetadata.DISTANCE.value());

// Remove all documents from the store
store.delete(documents().stream().map(doc -> doc.getId()).toList());

results = store.similaritySearch(SearchRequest.builder().query("Spring").topK(1).build());
assertThat(results).isEmpty();
}
});
}

@Test
void searchWithPartitionFilter() throws InterruptedException {
this.contextRunner.run(context -> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ public CassandraVectorStore store(CqlSession cqlSession, EmbeddingModel embeddin
.contentColumnName("body")
.embeddingColumnName("all_minilm_l6_v2_embedding")
.indexName("all_minilm_l6_v2_ann")
.disallowSchemaChanges(true)
.initializeSchema(false)
.addMetadataColumns(extraColumns)
.primaryKeyTranslator((List<Object> primaryKeys) -> {
// the deliminator used to join fields together into the document's id
Expand Down