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

Fix column creation, when adding additional normal and embedding colums), and make index name unique (for when there are multiple vector indexes in the same keyspace)

Browse Source 
And change stream to for-loop when converting List<Double> to Float[] for performance
This commit is contained in:
mck
2024-04-24 14:12:49 +02:00
committed by Christian Tzolov
parent a50969ec8b
commit daf131be2c
6 changed files with 40 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreProperties.java
View File

@@ -33,7 +33,7 @@ public class CassandraVectorStoreProperties {
private String table = CassandraVectorStoreConfig.DEFAULT_TABLE_NAME;
private String indexName = CassandraVectorStoreConfig.DEFAULT_INDEX_NAME;
private String indexName = null;
private String contentColumnName = CassandraVectorStoreConfig.DEFAULT_CONTENT_COLUMN_NAME;

2
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStorePropertiesTests.java
View File

@@ -34,7 +34,7 @@ class CassandraVectorStorePropertiesTests {
assertThat(props.getTable()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_TABLE_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.getIndexName()).isNull();
assertThat(props.getDisallowSchemaCreation()).isFalse();
assertThat(props.getFixedThreadPoolExecutorSize())
.isEqualTo(CassandraVectorStoreConfig.DEFAULT_ADD_CONCURRENCY);

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

@@ -206,7 +206,7 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean
@Override
public List<Document> similaritySearch(SearchRequest request) {
Preconditions.checkArgument(request.getTopK() <= 1000);
var embedding = this.embeddingClient.embed(request.getQuery()).stream().map(Double::floatValue).toList();
var embedding = toFloatArray(this.embeddingClient.embed(request.getQuery()));
CqlVector<Float> cqlVector = CqlVector.newInstance(embedding);
String whereClause = "";
@@ -350,4 +350,13 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean
return this.conf.primaryKeyTranslator.apply(primaryKeyValues);
}
private static Float[] toFloatArray(List<Double> embeddingDouble) {
Float[] embeddingFloat = new Float[embeddingDouble.size()];
int i = 0;
for (Double d : embeddingDouble) {
embeddingFloat[i++] = d.floatValue();
}
return embeddingFloat;
}
}

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

@@ -44,9 +44,12 @@ import com.datastax.oss.driver.api.querybuilder.schema.CreateTable;
import com.datastax.oss.driver.api.querybuilder.schema.CreateTableStart;
import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.lang.Nullable;
/**
* Configuration for the Cassandra vector store.
*
@@ -71,7 +74,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
public static final String DEFAULT_ID_NAME = "id";
public static final String DEFAULT_INDEX_NAME = "embedding_index";
public static final String DEFAULT_INDEX_SUFFIX = "idx";
public static final String DEFAULT_CONTENT_COLUMN_NAME = "content";
@@ -186,7 +189,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
private List<SchemaColumn> clusteringKeys = List.of();
private String indexName = DEFAULT_INDEX_NAME;
private String indexName = null;
private String contentColumnName = DEFAULT_CONTENT_COLUMN_NAME;
@@ -257,6 +260,8 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
return this;
}
/** defaults (if null) to '<table_name>_<embedding_column_name>_idx' **/
@Nullable
public Builder withIndexName(String name) {
this.indexName = name;
return this;
@@ -324,6 +329,9 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
}
public CassandraVectorStoreConfig build() {
if (null == this.indexName) {
this.indexName = String.format("%s_%s_%s", this.table, this.embeddingColumnName, DEFAULT_INDEX_SUFFIX);
}
for (SchemaColumn metadata : this.metadataColumns) {
Preconditions.checkArgument(
@@ -530,7 +538,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
// special case for embedding column, bc JAVA-3118, as above
StringBuilder alterTableStmt = new StringBuilder(((BuildableQuery) alterTable).asCql());
if (newColumns.isEmpty() && !addContent) {
alterTableStmt.append(" ADD ");
alterTableStmt.append(" ADD (");
}
else {
alterTableStmt.setLength(alterTableStmt.length() - 1);
@@ -539,7 +547,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
alterTableStmt.append(this.schema.embedding)
.append(" vector<float,")
.append(vectorDimension)
.append(">");
.append(">)");
logger.debug("Executing {}", alterTableStmt.toString());
this.session.execute(alterTableStmt.toString());

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

@@ -134,7 +134,8 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void ensureSchemaPartialCreation() {
this.contextRunner.run(context -> {
for (int i = 0; i < 4; ++i) {
int PARTIAL_FILES = 5;
for (int i = 0; i < PARTIAL_FILES; ++i) {
executeCqlFile(context, format("test_wiki_partial_%d_schema.cql", i));
var wrapper = createStore(context, List.of(), false, false);
try {
@@ -148,6 +149,10 @@ class CassandraRichSchemaVectorStoreIT {
wrapper.store().close();
}
}
// make sure there's not more files to test
Assertions.assertThrows(IOException.class, () -> {
executeCqlFile(context, format("test_wiki_partial_%d_schema.cql", PARTIAL_FILES));
});
});
}

10
vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_4_schema.cql Normal file
View File

@@ -0,0 +1,10 @@
CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1};
CREATE TABLE IF NOT EXISTS test_wikidata.articles (
wiki text,
language text,
title text,
chunk_no int,
messages text,
PRIMARY KEY ((wiki, language, title), chunk_no)
);