Skip to content

Cassandra vector store builder refactoring #1912

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed

Cassandra vector store builder refactoring #1912

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
228056e
Cassandra vector store builder refactoring
sobychacko Dec 12, 2024
92b420c
Removing the usage of the config class
sobychacko Dec 12, 2024
782ec0d
Addressing PR review
sobychacko Dec 14, 2024
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
4 changes: 2 additions & 2 deletions ...ramework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryAutoConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@

import com.datastax.oss.driver.api.core.CqlSession;

import org.springframework.ai.chat.memory.CassandraChatMemory;
import org.springframework.ai.chat.memory.CassandraChatMemoryConfig;
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemory;
Copy link
Contributor

@michaelsembwever michaelsembwever Dec 12, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

are all vectorstore modules moving their classes to subpackages ? (i don't see that yet with pgvector. has there been some discussion/agreement to this somewhere we can link to ?)

Copy link
Contributor Author

@sobychacko sobychacko Dec 14, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, we are in the process of moving the vector store classes under that structure - in this case - org.springframework.ai.chat.memory.cassandra and springframework.ai.vectorstore.cassandra. The issue was that the vector store classes currently use the same package name in all the different modules - springframework.ai.vectorstore.

Copy link
Contributor

@michaelsembwever michaelsembwever Dec 14, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

was hoping that there was some link you could provide to where the discussion had been had and landed on that conclusion and action. but it's ok, i trust.

Copy link
Member

@markpollack markpollack Dec 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

next time we will reach out to you, i wish we had a more centralized way to communicate changes to frequent contributors. we are working on this.

import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryConfig;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.cassandra.CassandraAutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
Expand Down
2 changes: 1 addition & 1 deletion ...springframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryProperties.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import org.slf4j.LoggerFactory;

import org.springframework.ai.autoconfigure.chat.memory.CommonChatMemoryProperties;
import org.springframework.ai.chat.memory.CassandraChatMemoryConfig;
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryConfig;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.lang.Nullable;

Expand Down
37 changes: 16 additions & 21 deletions ...amework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,7 @@
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.vectorstore.CassandraVectorStore;
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig;
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.autoconfigure.AutoConfiguration;
Expand Down Expand Up @@ -63,25 +62,21 @@ public CassandraVectorStore vectorStore(EmbeddingModel embeddingModel, Cassandra
ObjectProvider<VectorStoreObservationConvention> customObservationConvention,
BatchingStrategy batchingStrategy) {

var builder = CassandraVectorStoreConfig.builder().withCqlSession(cqlSession);

builder = builder.withKeyspaceName(properties.getKeyspace())
.withTableName(properties.getTable())
.withContentColumnName(properties.getContentColumnName())
.withEmbeddingColumnName(properties.getEmbeddingColumnName())
.withIndexName(properties.getIndexName())
.withFixedThreadPoolExecutorSize(properties.getFixedThreadPoolExecutorSize());

if (!properties.isInitializeSchema()) {
builder = builder.disallowSchemaChanges();
}
if (properties.getReturnEmbeddings()) {
builder = builder.returnEmbeddings();
}

return new CassandraVectorStore(builder.build(), embeddingModel,
observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
customObservationConvention.getIfAvailable(() -> null), batchingStrategy);
return CassandraVectorStore.builder()
.session(cqlSession)
.keyspace(properties.getKeyspace())
.table(properties.getTable())
.contentColumnName(properties.getContentColumnName())
.embeddingColumnName(properties.getEmbeddingColumnName())
.indexName(properties.getIndexName())
.fixedThreadPoolExecutorSize(properties.getFixedThreadPoolExecutorSize())
.disallowSchemaChanges(!properties.isInitializeSchema())
.returnEmbeddings(properties.getReturnEmbeddings())
Comment on lines +73 to +74
Copy link
Contributor

@michaelsembwever michaelsembwever Dec 12, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

there's a behaviour change here, in that we're no longer respecting the builder's default, imposing the autoConfiguration's default over it. (the builder's default should be canonical)

i'm not sure how important this is… 🤷

Copy link
Member

@markpollack markpollack Dec 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, we should be consistent. I'll create a new issue to capture this and discuss.

.embeddingModel(embeddingModel)
.observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
.customObservationConvention(customObservationConvention.getIfAvailable(() -> null))
.batchingStrategy(batchingStrategy)
.build();
}

@Bean
Expand Down
12 changes: 6 additions & 6 deletions ...pringframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreProperties.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
import org.slf4j.LoggerFactory;

import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig;
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore;
import org.springframework.boot.context.properties.ConfigurationProperties;

/**
Expand All @@ -37,19 +37,19 @@ public class CassandraVectorStoreProperties extends CommonVectorStoreProperties

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

private String keyspace = CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME;
private String keyspace = CassandraVectorStore.DEFAULT_KEYSPACE_NAME;

private String table = CassandraVectorStoreConfig.DEFAULT_TABLE_NAME;
private String table = CassandraVectorStore.DEFAULT_TABLE_NAME;

private String indexName = null;

private String contentColumnName = CassandraVectorStoreConfig.DEFAULT_CONTENT_COLUMN_NAME;
private String contentColumnName = CassandraVectorStore.DEFAULT_CONTENT_COLUMN_NAME;

private String embeddingColumnName = CassandraVectorStoreConfig.DEFAULT_EMBEDDING_COLUMN_NAME;
private String embeddingColumnName = CassandraVectorStore.DEFAULT_EMBEDDING_COLUMN_NAME;

private boolean returnEmbeddings = false;

private int fixedThreadPoolExecutorSize = CassandraVectorStoreConfig.DEFAULT_ADD_CONCURRENCY;
private int fixedThreadPoolExecutorSize = CassandraVectorStore.DEFAULT_ADD_CONCURRENCY;

public String getKeyspace() {
return this.keyspace;
Expand Down
2 changes: 1 addition & 1 deletion ...mework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryAutoConfigurationIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.utility.DockerImageName;

import org.springframework.ai.chat.memory.CassandraChatMemory;
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
Expand Down
2 changes: 1 addition & 1 deletion ...ngframework/ai/autoconfigure/chat/memory/cassandra/CassandraChatMemoryPropertiesTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@

import org.junit.jupiter.api.Test;

import org.springframework.ai.chat.memory.CassandraChatMemoryConfig;
import org.springframework.ai.chat.memory.cassandra.CassandraChatMemoryConfig;

import static org.assertj.core.api.Assertions.assertThat;

Expand Down
14 changes: 7 additions & 7 deletions ...framework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStorePropertiesTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,8 @@

import org.junit.jupiter.api.Test;

import org.springframework.ai.vectorstore.CassandraVectorStoreConfig;
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore;
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStoreConfig;

import static org.assertj.core.api.Assertions.assertThat;

Expand All @@ -31,13 +32,12 @@ class CassandraVectorStorePropertiesTests {
@Test
void defaultValues() {
var props = new CassandraVectorStoreProperties();
assertThat(props.getKeyspace()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME);
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.getKeyspace()).isEqualTo(CassandraVectorStore.DEFAULT_KEYSPACE_NAME);
assertThat(props.getTable()).isEqualTo(CassandraVectorStore.DEFAULT_TABLE_NAME);
assertThat(props.getContentColumnName()).isEqualTo(CassandraVectorStore.DEFAULT_CONTENT_COLUMN_NAME);
assertThat(props.getEmbeddingColumnName()).isEqualTo(CassandraVectorStore.DEFAULT_EMBEDDING_COLUMN_NAME);
assertThat(props.getIndexName()).isNull();
assertThat(props.getFixedThreadPoolExecutorSize())
.isEqualTo(CassandraVectorStoreConfig.DEFAULT_ADD_CONCURRENCY);
assertThat(props.getFixedThreadPoolExecutorSize()).isEqualTo(CassandraVectorStore.DEFAULT_ADD_CONCURRENCY);
}

@Test
Expand Down
7 changes: 4 additions & 3 deletions ...k/ai/chat/memory/CassandraChatMemory.java → ...memory/cassandra/CassandraChatMemory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -32,7 +32,8 @@
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;
Expand All @@ -42,7 +43,7 @@
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
Expand Down
2 changes: 1 addition & 1 deletion ...hat/memory/CassandraChatMemoryConfig.java → .../cassandra/CassandraChatMemoryConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading
Loading