Skip to content

refactor ChromaVectorStore builder API #1902

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

refactor ChromaVectorStore builder API #1902

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
50a9a92
refactor ChromaVectorStore builder API
sobychacko Dec 10, 2024
ad4fcd9
Addressing PR review
sobychacko Dec 10, 2024
1e816a5
cleanup
sobychacko Dec 10, 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
3 changes: 2 additions & 1 deletion ...ringframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,8 @@ public ChromaVectorStore vectorStore(EmbeddingModel embeddingModel, ChromaApi ch
ChromaVectorStoreProperties storeProperties, ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<VectorStoreObservationConvention> customObservationConvention,
BatchingStrategy chromaBatchingStrategy) {
return ChromaVectorStore.builder(chromaApi)
return ChromaVectorStore.builder()
.chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName(storeProperties.getCollectionName())
.initializeSchema(storeProperties.isInitializeSchema())
Expand Down
17 changes: 11 additions & 6 deletions ...roma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,8 @@ public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, Str
boolean initializeSchema, ObservationRegistry observationRegistry,
VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) {

this(builder(chromaApi).embeddingModel(embeddingModel)
this(builder().chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName(collectionName)
.initializeSchema(initializeSchema)
.observationRegistry(observationRegistry)
Expand All @@ -113,8 +114,11 @@ public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, Str
/**
* @param builder {@link Builder} for chroma vector store
*/
private ChromaVectorStore(ChromaBuilder builder) {
protected ChromaVectorStore(ChromaBuilder builder) {
super(builder);

Assert.notNull(builder.chromaApi, "ChromaApi must not be null");

this.chromaApi = builder.chromaApi;
this.collectionName = builder.collectionName;
this.initializeSchema = builder.initializeSchema;
Expand Down Expand Up @@ -151,8 +155,8 @@ public void afterPropertiesSet() throws Exception {
}
}

public static ChromaBuilder builder(ChromaApi chromaApi) {
return new ChromaBuilder(chromaApi);
public static ChromaBuilder builder() {
return new ChromaBuilder();
}

@Override
Expand Down Expand Up @@ -274,7 +278,7 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str

public static class ChromaBuilder extends AbstractVectorStoreBuilder<ChromaBuilder> {

private final ChromaApi chromaApi;
private ChromaApi chromaApi;

private String collectionName = DEFAULT_COLLECTION_NAME;

Expand All @@ -286,9 +290,10 @@ public static class ChromaBuilder extends AbstractVectorStoreBuilder<ChromaBuild

private boolean initializeImmediately = false;

public ChromaBuilder(ChromaApi chromaApi) {
public ChromaBuilder chromaApi(ChromaApi chromaApi) {
Assert.notNull(chromaApi, "ChromaApi must not be null");
this.chromaApi = chromaApi;
return this;
}

/**
Expand Down
3 changes: 2 additions & 1 deletion ...store/src/test/java/org/springframework/ai/chroma/vectorstore/BasicAuthChromaWhereIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,8 @@ public ChromaApi chromaApi(RestClient.Builder builder) {

@Bean
public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi) {
return ChromaVectorStore.builder(chromaApi)
return ChromaVectorStore.builder()
.chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName("TestCollection")
.initializeSchema(true)
Expand Down
8 changes: 5 additions & 3 deletions ...-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,8 @@ void shouldUseExistingCollectionWhenSchemaInitializationDisabled() { // initiali
assertThat(collection).isNotNull();
assertThat(collection.name()).isEqualTo("test-collection");

ChromaVectorStore store = ChromaVectorStore.builder(this.chromaApi)
ChromaVectorStore store = ChromaVectorStore.builder()
.chromaApi(this.chromaApi)
.embeddingModel(this.embeddingModel)
.collectionName("test-collection")
.initializeImmediately(true)
Expand All @@ -219,7 +220,7 @@ void shouldUseExistingCollectionWhenSchemaInitializationDisabled() { // initiali

@Test
void shouldCreateNewCollectionWhenSchemaInitializationEnabled() {
ChromaVectorStore store = new ChromaVectorStore.ChromaBuilder(this.chromaApi)
ChromaVectorStore store = new ChromaVectorStore.ChromaBuilder().chromaApi(this.chromaApi)
.embeddingModel(this.embeddingModel)
.collectionName("new-collection")
.initializeSchema(true)
Expand All @@ -236,7 +237,8 @@ void shouldCreateNewCollectionWhenSchemaInitializationEnabled() {

@Test
void shouldFailWhenCollectionDoesNotExist() {
assertThatThrownBy(() -> new ChromaVectorStore.ChromaBuilder(this.chromaApi).embeddingModel(this.embeddingModel)
assertThatThrownBy(() -> new ChromaVectorStore.ChromaBuilder().chromaApi(this.chromaApi)
.embeddingModel(this.embeddingModel)
.collectionName("non-existent")
.initializeSchema(false)
.initializeImmediately(true)
Expand Down
3 changes: 2 additions & 1 deletion ...ma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -252,7 +252,8 @@ public ChromaApi chromaApi(RestClient.Builder builder) {

@Bean
public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi) {
return ChromaVectorStore.builder(chromaApi)
return ChromaVectorStore.builder()
.chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName("TestCollection")
.initializeSchema(true)
Expand Down
3 changes: 2 additions & 1 deletion ...c/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreObservationIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,8 @@ public ChromaApi chromaApi(RestClient.Builder builder) {
@Bean
public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi,
ObservationRegistry observationRegistry) {
return ChromaVectorStore.builder(chromaApi)
return ChromaVectorStore.builder()
.chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName("TestCollection")
.initializeSchema(true)
Expand Down
3 changes: 2 additions & 1 deletion ...re/src/test/java/org/springframework/ai/chroma/vectorstore/TokenSecuredChromaWhereIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,8 @@ public ChromaApi chromaApi(RestClient.Builder builder) {

@Bean
public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi) {
return ChromaVectorStore.builder(chromaApi)
return ChromaVectorStore.builder()
.chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName("TestCollection")
.initializeSchema(true)
Expand Down
Loading